From 5d231c64cd04716f858a6e35df02f3a7eeae1f53 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 3 Jun 2026 04:40:22 -0700 Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub. PiperOrigin-RevId: 925939529 Change-Id: Idb92c2d50d5be9df7b6e91113b576e9c706847ca --- .../mjx/third_party/mujoco_warp/_src/io.py | 12 +- .../third_party/mujoco_warp/_src/render.py | 770 ++++++++---------- .../third_party/mujoco_warp/_src/support.py | 6 +- .../mjx/third_party/mujoco_warp/_src/types.py | 6 + mjx/mujoco/mjx/warp/bvh.py | 11 +- mjx/mujoco/mjx/warp/collision_driver.py | 8 +- mjx/mujoco/mjx/warp/forward.py | 8 +- mjx/mujoco/mjx/warp/render.py | 13 +- mjx/mujoco/mjx/warp/smooth.py | 8 +- mjx/mujoco/mjx/warp/types.py | 21 +- 10 files changed, 379 insertions(+), 484 deletions(-) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py index ffd9f8f3..2810bc79 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py @@ -3058,11 +3058,11 @@ def create_render_context( # Locate skybox texture skybox_tex_ids = np.nonzero(mjm.tex_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX)[0] if mjm.ntex else np.array([], dtype=int) - if render_skybox: - assert skybox_tex_ids.size > 0, "render_skybox=True but the model has no texture with type mjTEXTURE_SKYBOX" + if render_skybox and skybox_tex_ids.size > 0: skybox_tex_id = int(skybox_tex_ids[0]) skybox_face_width = int(mjm.tex_width[skybox_tex_id]) else: + render_skybox = False skybox_tex_id = -1 skybox_face_width = 1 @@ -3157,6 +3157,13 @@ def create_render_context( bvh_ngeom = len(geom_enabled_idx) + # Geom types present among enabled geoms, plus FLEX when flex primitives exist. + # Used to statically eliminate unused intersection branches in the ray-cast kernels. + geom_ray_types = set(int(t) for t in mjm.geom_type[geom_enabled_idx]) + if len(flex_geom_flexid) > 0: + geom_ray_types.add(int(types.GeomType.FLEX)) + geom_ray_types = tuple(sorted(geom_ray_types)) + rc = types.RenderContext( nrender=ncam, cam_res=cam_res_arr, @@ -3210,6 +3217,7 @@ def create_render_context( znear=znear, total_rays=int(total), enable_backface_culling=enable_backface_culling, + geom_ray_types=geom_ray_types, ) bvh.build_scene_bvh(mjm, mjd, rc, nworld) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py index 7afb140b..05780e5b 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py @@ -151,449 +151,346 @@ def sample_skybox( return wp.vec3(color[0], color[1], color[2]) -# TODO: Investigate combining cast_ray and cast_ray_first_hit -@wp.func -def cast_ray( - # Model: - geom_type: wp.array[int], - geom_dataid: wp.array2d[int], - geom_size: wp.array2d[wp.vec3], - flex_vertadr: wp.array[int], - flex_edge: wp.array[wp.vec2i], - flex_radius: wp.array[float], - # Data in: - geom_xpos_in: wp.array2d[wp.vec3], - geom_xmat_in: wp.array2d[wp.mat33], - flexvert_xpos_in: wp.array2d[wp.vec3], - # In: - bvh_id: wp.uint64, - group_root: int, - worldid: int, - bvh_ngeom: int, - flex_bvh_ngeom: int, - enabled_geom_ids: wp.array[int], - mesh_bvh_id: wp.array[wp.uint64], - hfield_bvh_id: wp.array[wp.uint64], - flex_geom_flexid: wp.array[int], - flex_geom_edgeid: wp.array[int], - flex_bvh_id: wp.array[wp.uint64], - flex_group_root: wp.array2d[int], - ray_origin_world: wp.vec3, - ray_dir_world: wp.vec3, - cull_backfaces: bool, -) -> Tuple[int, float, wp.vec3, float, float, int, int]: - dist = float(MJ_MAXVAL) - normal = wp.vec3(0.0, 0.0, 0.0) - geom_id = int(-1) - bary_u = float(0.0) - bary_v = float(0.0) - face_idx = int(-1) - geom_mesh_id = int(-1) +def _make_cast_ray(geom_ray_types: Tuple[int], first_hit: bool = False) -> wp.Function: + """Build a ray-cast func specialized to the geom types present in the scene. - query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root) - bounds_nr = int(0) - ngeom = bvh_ngeom + flex_bvh_ngeom + geom_ray_types is the set of GeomType int values that actually occur, so the + per-type intersection branches for absent types are eliminated at compile time + via wp.static, avoiding the register pressure of unreachable code paths. - while wp.bvh_query_next(query, bounds_nr, dist): - gi_global = bounds_nr - local_id = gi_global - (worldid * ngeom) + first_hit selects the variant (also resolved at compile time via wp.static): + - False: full closest-hit cast. Returns the closest hit's full surface data. + - True: any-hit cast (shadow rays). Uses the cheaper any-hit mesh/flex + intersections and returns on the first hit within max_dist. The result is + still the full tuple; callers test geom_id != -1 to detect a hit. + """ - d = float(-1.0) - hit_mesh_id = int(-1) - u = float(0.0) - v = float(0.0) - f = int(-1) - n = wp.vec3(0.0, 0.0, 0.0) - hit_geom_id = int(-1) + @wp.func + def cast_ray( + # Model: + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + flex_vertadr: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], + # Data in: + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], + # In: + bvh_id: wp.uint64, + group_root: int, + worldid: int, + bvh_ngeom: int, + flex_bvh_ngeom: int, + enabled_geom_ids: wp.array[int], + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], + flex_bvh_id: wp.array[wp.uint64], + flex_group_root: wp.array2d[int], + ray_origin_world: wp.vec3, + ray_dir_world: wp.vec3, + max_dist: float, + cull_backfaces: bool, + ) -> Tuple[int, float, wp.vec3, float, float, int, int]: + dist = max_dist + normal = wp.vec3(0.0, 0.0, 0.0) + geom_id = int(-1) + bary_u = float(0.0) + bary_v = float(0.0) + face_idx = int(-1) + geom_mesh_id = int(-1) - if local_id < bvh_ngeom: - gi = enabled_geom_ids[local_id] - gtype = geom_type[gi] - else: - gi = local_id - bvh_ngeom - gtype = GeomType.FLEX + query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root) + bounds_nr = int(0) + ngeom = bvh_ngeom + flex_bvh_ngeom - hit_geom_id = gi + while wp.bvh_query_next(query, bounds_nr, dist): + gi_global = bounds_nr + local_id = gi_global - (worldid * ngeom) - # TODO: Investigate branch elimination with static loop unrolling - 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], - ray_origin_world, - ray_dir_world, - ) - if gtype == GeomType.HFIELD: - d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( - hfield_bvh_id, - geom_dataid[worldid % geom_dataid.shape[0], gi], - geom_xpos_in[worldid, gi], - geom_xmat_in[worldid, gi], - ray_origin_world, - ray_dir_world, - dist, - cull_backfaces, - ) - if gtype == GeomType.SPHERE: - d, n = ray_sphere( - geom_xpos_in[worldid, gi], - geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0], - ray_origin_world, - ray_dir_world, - ) - 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], - ray_origin_world, - ray_dir_world, - ) - 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], - ray_origin_world, - ray_dir_world, - ) - 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], - ray_origin_world, - ray_dir_world, - ) - if gtype == GeomType.BOX: - d, all, n = ray_box( - 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 gtype == GeomType.MESH: - d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh( - mesh_bvh_id, - geom_dataid[worldid % geom_dataid.shape[0], gi], - geom_xpos_in[worldid, gi], - geom_xmat_in[worldid, gi], - ray_origin_world, - ray_dir_world, - dist, - cull_backfaces, - ) - if gtype == GeomType.FLEX: - hit_geom_id = -2 - flexid = flex_geom_flexid[gi] - edge_id = flex_geom_edgeid[gi] + d = float(-1.0) + hit_mesh_id = int(-1) + u = float(0.0) + v = float(0.0) + f = int(-1) + n = wp.vec3(0.0, 0.0, 0.0) + hit_geom_id = int(-1) - if edge_id >= 0: - edge = flex_edge[edge_id] - vert_adr = flex_vertadr[flexid] - v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]] - v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]] - pos = 0.5 * (v0 + v1) - vec = v1 - v0 - - length = wp.length(vec) - edgeq = math.quat_z2vec(vec) - mat = math.quat_to_mat(edgeq) - size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0) - - d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world) - hit_mesh_id = flexid + if local_id < bvh_ngeom: + gi = enabled_geom_ids[local_id] + gtype = geom_type[gi] else: - flex_gr = flex_group_root[worldid, flexid] - d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist) - if d >= 0.0: - hit_mesh_id = flexid + gi = local_id - bvh_ngeom + gtype = GeomType.FLEX - # Backface cull: drop exit-face hits when the ray origin is inside the geom, - # matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule. - if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0: - d = -1.0 + hit_geom_id = gi - if d >= 0.0 and d < dist: - dist = d - normal = n - geom_id = hit_geom_id - bary_u = u - bary_v = v - face_idx = f - geom_mesh_id = hit_mesh_id + 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], + ray_origin_world, + ray_dir_world, + ) + if wp.static(int(GeomType.HFIELD) in geom_ray_types): + if gtype == GeomType.HFIELD: + d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( + hfield_bvh_id, + geom_dataid[worldid % geom_dataid.shape[0], gi], + geom_xpos_in[worldid, gi], + geom_xmat_in[worldid, gi], + ray_origin_world, + ray_dir_world, + dist, + cull_backfaces, + ) + if wp.static(int(GeomType.SPHERE) in geom_ray_types): + if gtype == GeomType.SPHERE: + d, n = ray_sphere( + geom_xpos_in[worldid, gi], + geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0], + ray_origin_world, + ray_dir_world, + ) + 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], + ray_origin_world, + ray_dir_world, + ) + 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], + ray_origin_world, + ray_dir_world, + ) + 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], + ray_origin_world, + ray_dir_world, + ) + if wp.static(int(GeomType.BOX) in geom_ray_types): + if gtype == GeomType.BOX: + d, all, n = ray_box( + 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 wp.static(int(GeomType.MESH) in geom_ray_types): + if gtype == GeomType.MESH: + if wp.static(first_hit): + hit = ray_mesh_with_bvh_anyhit( + mesh_bvh_id, + geom_dataid[worldid % geom_dataid.shape[0], gi], + geom_xpos_in[worldid, gi], + geom_xmat_in[worldid, gi], + ray_origin_world, + ray_dir_world, + dist, + ) + d = 0.0 if hit else -1.0 + else: + d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh( + mesh_bvh_id, + geom_dataid[worldid % geom_dataid.shape[0], gi], + geom_xpos_in[worldid, gi], + geom_xmat_in[worldid, gi], + ray_origin_world, + ray_dir_world, + dist, + cull_backfaces, + ) + if wp.static(int(GeomType.FLEX) in geom_ray_types): + if gtype == GeomType.FLEX: + hit_geom_id = -2 + flexid = flex_geom_flexid[gi] + edge_id = flex_geom_edgeid[gi] - return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id + if edge_id >= 0: + edge = flex_edge[edge_id] + vert_adr = flex_vertadr[flexid] + v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]] + v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]] + pos = 0.5 * (v0 + v1) + vec = v1 - v0 + length = wp.length(vec) + edgeq = math.quat_z2vec(vec) + mat = math.quat_to_mat(edgeq) + size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0) -@wp.func -def cast_ray_first_hit( - # Model: - geom_type: wp.array[int], - geom_dataid: wp.array2d[int], - geom_size: wp.array2d[wp.vec3], - flex_vertadr: wp.array[int], - flex_edge: wp.array[wp.vec2i], - flex_radius: wp.array[float], - # Data in: - geom_xpos_in: wp.array2d[wp.vec3], - geom_xmat_in: wp.array2d[wp.mat33], - flexvert_xpos_in: wp.array2d[wp.vec3], - # In: - bvh_id: wp.uint64, - group_root: int, - worldid: int, - bvh_ngeom: int, - bvh_nflexgeom: int, - enabled_geom_ids: wp.array[int], - mesh_bvh_id: wp.array[wp.uint64], - hfield_bvh_id: wp.array[wp.uint64], - flex_geom_flexid: wp.array[int], - flex_geom_edgeid: wp.array[int], - flex_bvh_id: wp.array[wp.uint64], - flex_group_root: wp.array2d[int], - ray_origin_world: wp.vec3, - ray_dir_world: wp.vec3, - max_dist: float, - cull_backfaces: bool, -) -> bool: - """A simpler version of casting rays that only checks for the first hit.""" - query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root) - bounds_nr = int(0) - ngeom = bvh_ngeom + bvh_nflexgeom + d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world) + hit_mesh_id = flexid + else: + if wp.static(first_hit): + hit = ray_flex_with_bvh_anyhit( + flex_bvh_id, + flexid, + flex_group_root[worldid, flexid], + ray_origin_world, + ray_dir_world, + dist, + ) + d = 0.0 if hit else -1.0 + else: + flex_gr = flex_group_root[worldid, flexid] + d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist) + if d >= 0.0: + hit_mesh_id = flexid - while wp.bvh_query_next(query, bounds_nr, max_dist): - gi_global = bounds_nr - local_id = gi_global - (worldid * ngeom) + # Backface cull: drop exit-face hits when the ray origin is inside the geom, + # matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule. Strict `> 0` keeps + # tangent hits and skips branches with a zero-vector normal (any-hit). + if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0: + d = -1.0 - d = float(-1.0) - n = wp.vec3(0.0, 0.0, 0.0) - - if local_id < bvh_ngeom: - gi = enabled_geom_ids[local_id] - gtype = geom_type[gi] - else: - gi = local_id - bvh_ngeom - gtype = GeomType.FLEX - - # TODO: Investigate branch elimination with static loop unrolling - 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], - ray_origin_world, - ray_dir_world, - ) - if gtype == GeomType.HFIELD: - d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh( - hfield_bvh_id, - geom_dataid[worldid % geom_dataid.shape[0], gi], - geom_xpos_in[worldid, gi], - geom_xmat_in[worldid, gi], - ray_origin_world, - ray_dir_world, - max_dist, - cull_backfaces, - ) - if gtype == GeomType.SPHERE: - d, n = ray_sphere( - geom_xpos_in[worldid, gi], - geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0], - ray_origin_world, - ray_dir_world, - ) - 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], - ray_origin_world, - ray_dir_world, - ) - 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], - ray_origin_world, - ray_dir_world, - ) - 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], - ray_origin_world, - ray_dir_world, - ) - if gtype == GeomType.BOX: - d, all, n = ray_box( - 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 gtype == GeomType.MESH: - hit = ray_mesh_with_bvh_anyhit( - mesh_bvh_id, - geom_dataid[worldid % geom_dataid.shape[0], gi], - geom_xpos_in[worldid, gi], - geom_xmat_in[worldid, gi], - ray_origin_world, - ray_dir_world, - max_dist, - ) - d = 0.0 if hit else -1.0 - if gtype == GeomType.FLEX: - flexid = flex_geom_flexid[gi] - edge_id = flex_geom_edgeid[gi] - - if edge_id >= 0: - edge = flex_edge[edge_id] - vert_adr = flex_vertadr[flexid] - v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]] - v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]] - pos = 0.5 * (v0 + v1) - vec = v1 - v0 - - length = wp.length(vec) - edgeq = math.quat_z2vec(vec) - mat = math.quat_to_mat(edgeq) - size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0) - - d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world) + if wp.static(first_hit): + # Any-hit: return as soon as anything is in range; surface data is unused. + if d >= 0.0 and d < dist: + return hit_geom_id, d, n, u, v, f, hit_mesh_id else: - hit = ray_flex_with_bvh_anyhit( - flex_bvh_id, - flexid, - flex_group_root[worldid, flexid], - ray_origin_world, - ray_dir_world, - max_dist, - ) - d = 0.0 if hit else -1.0 + if d >= 0.0 and d < dist: + dist = d + normal = n + geom_id = hit_geom_id + bary_u = u + bary_v = v + face_idx = f + geom_mesh_id = hit_mesh_id - # Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent - # hits and skips branches with a zero-vector normal (mesh/flex anyhit). - if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0: - d = -1.0 + return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id - if d >= 0.0 and d < max_dist: - return True - - return False + return cast_ray -@wp.func -def compute_lighting( - # Model: - geom_type: wp.array[int], - geom_dataid: wp.array2d[int], - geom_size: wp.array2d[wp.vec3], - flex_vertadr: wp.array[int], - flex_edge: wp.array[wp.vec2i], - flex_radius: wp.array[float], - # Data in: - geom_xpos_in: wp.array2d[wp.vec3], - geom_xmat_in: wp.array2d[wp.mat33], - flexvert_xpos_in: wp.array2d[wp.vec3], - # In: - use_shadows: bool, - bvh_id: wp.uint64, - group_root: int, - bvh_ngeom: int, - bvh_nflexgeom: int, - enabled_geom_ids: wp.array[int], - worldid: int, - mesh_bvh_id: wp.array[wp.uint64], - hfield_bvh_id: wp.array[wp.uint64], - flex_geom_flexid: wp.array[int], - flex_geom_edgeid: wp.array[int], - flex_bvh_id: wp.array[wp.uint64], - flex_group_root: wp.array2d[int], - lightactive: bool, - lighttype: int, - lightcastshadow: bool, - lightpos: wp.vec3, - lightdir: wp.vec3, - normal: wp.vec3, - hitpoint: wp.vec3, - cull_backfaces: bool, -) -> float: - light_contribution = float(0.0) +def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function: + """Build specialized compute_lighting.""" - # TODO: We should probably only be looping over active lights - # in the first place with a static loop of enabled light idx? - if not lightactive: - return light_contribution + @wp.func + def compute_lighting( + # Model: + geom_type: wp.array[int], + geom_dataid: wp.array2d[int], + geom_size: wp.array2d[wp.vec3], + flex_vertadr: wp.array[int], + flex_edge: wp.array[wp.vec2i], + flex_radius: wp.array[float], + # Data in: + geom_xpos_in: wp.array2d[wp.vec3], + geom_xmat_in: wp.array2d[wp.mat33], + flexvert_xpos_in: wp.array2d[wp.vec3], + # In: + use_shadows: bool, + bvh_id: wp.uint64, + group_root: int, + bvh_ngeom: int, + bvh_nflexgeom: int, + enabled_geom_ids: wp.array[int], + worldid: int, + mesh_bvh_id: wp.array[wp.uint64], + hfield_bvh_id: wp.array[wp.uint64], + flex_geom_flexid: wp.array[int], + flex_geom_edgeid: wp.array[int], + flex_bvh_id: wp.array[wp.uint64], + flex_group_root: wp.array2d[int], + lightactive: bool, + lighttype: int, + lightcastshadow: bool, + lightpos: wp.vec3, + lightdir: wp.vec3, + normal: wp.vec3, + hitpoint: wp.vec3, + cull_backfaces: bool, + ) -> float: + light_contribution = float(0.0) - L = wp.vec3(0.0, 0.0, 0.0) - dist_to_light = float(MJ_MAXVAL) - attenuation = float(1.0) + # TODO: We should probably only be looping over active lights + # in the first place with a static loop of enabled light idx? + if not lightactive: + return light_contribution - if lighttype == 1: # directional light - L = wp.normalize(-lightdir) - else: - L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint) - attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light) - if lighttype == 0: # spot light - spot_dir = wp.normalize(lightdir) - cos_theta = wp.dot(-L, spot_dir) - spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) / (0.95 - 0.85))) - attenuation = attenuation * spot_factor + L = wp.vec3(0.0, 0.0, 0.0) + dist_to_light = float(MJ_MAXVAL) + attenuation = float(1.0) - ndotl = wp.max(0.0, wp.dot(normal, L)) - if ndotl == 0.0: - return light_contribution - - visible = float(1.0) - - if use_shadows and lightcastshadow: - # Nudge the origin slightly along the surface normal to avoid - # self-intersection when casting shadow rays - eps = 1.0e-4 - shadow_origin = hitpoint + normal * eps - # Distance-limited shadows: cap by dist_to_light (for non-directional) - max_t = float(dist_to_light - 1.0e-3) if lighttype == 1: # directional light - max_t = float(1.0e8) + L = wp.normalize(-lightdir) + else: + L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint) + attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light) + if lighttype == 0: # spot light + spot_dir = wp.normalize(lightdir) + cos_theta = wp.dot(-L, spot_dir) + spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) * 10.0)) + attenuation = attenuation * spot_factor - shadow_hit = cast_ray_first_hit( - geom_type, - geom_dataid, - geom_size, - flex_vertadr, - flex_edge, - flex_radius, - geom_xpos_in, - geom_xmat_in, - flexvert_xpos_in, - bvh_id, - group_root, - worldid, - bvh_ngeom, - bvh_nflexgeom, - enabled_geom_ids, - mesh_bvh_id, - hfield_bvh_id, - flex_geom_flexid, - flex_geom_edgeid, - flex_bvh_id, - flex_group_root, - shadow_origin, - L, - max_t, - cull_backfaces, - ) + ndotl = wp.max(0.0, wp.dot(normal, L)) + if ndotl == 0.0: + return light_contribution - if shadow_hit: - visible = 0.3 + visible = float(1.0) - return ndotl * attenuation * visible + if use_shadows and lightcastshadow: + # Nudge the origin slightly along the surface normal to avoid + # self-intersection when casting shadow rays + shadow_origin = hitpoint + normal * 1.0e-4 + # Distance-limited shadows: cap by dist_to_light (for non-directional) + max_t = dist_to_light - 1.0e-3 + if lighttype == 1: # directional light + max_t = 1.0e8 + + shadow_geom_id, shadow_d, shadow_n, shadow_u, shadow_v, shadow_f, shadow_mesh_id = cast_ray_first_hit( + geom_type, + geom_dataid, + geom_size, + flex_vertadr, + flex_edge, + flex_radius, + geom_xpos_in, + geom_xmat_in, + flexvert_xpos_in, + bvh_id, + group_root, + worldid, + bvh_ngeom, + bvh_nflexgeom, + enabled_geom_ids, + mesh_bvh_id, + hfield_bvh_id, + flex_geom_flexid, + flex_geom_edgeid, + flex_bvh_id, + flex_group_root, + shadow_origin, + L, + max_t, + cull_backfaces, + ) + + if shadow_geom_id != -1: + visible = 0.3 + + return ndotl * attenuation * visible + + return compute_lighting @event_scope @@ -611,6 +508,13 @@ def render(m: Model, d: Data, rc: RenderContext): rc.depth_data.fill_(0.0) rc.seg_data.fill_(wp.vec2i(-1, -1)) + # Specialize the ray-cast helpers to the geom types present in the scene so the + # compiler eliminates intersection branches for absent types. + geom_ray_types = rc.geom_ray_types + cast_ray = _make_cast_ray(geom_ray_types, first_hit=False) + cast_ray_first_hit = _make_cast_ray(geom_ray_types, first_hit=True) + compute_lighting = _make_compute_lighting(cast_ray_first_hit) + @wp.kernel(module="unique", enable_backward=False) def _render_megakernel( # Model: @@ -676,31 +580,31 @@ def render(m: Model, d: Data, rc: RenderContext): ): worldid, rayid = wp.tid() - # Map global rayid -> (cam_idx, rayid_local) using cumulative sizes - cam_idx = int(-1) + # Map global rayid -> (camid, rayid_local) using cumulative sizes + camid = int(-1) rayid_local = int(-1) accum = int(0) for i in range(nrender): num_i = cam_res[i][0] * cam_res[i][1] if rayid < accum + num_i: - cam_idx = i + camid = i rayid_local = rayid - accum break accum += num_i - if cam_idx == -1 or rayid_local < 0: + if camid == -1 or rayid_local < 0: return - if not render_rgb[cam_idx] and not render_depth[cam_idx] and not render_seg[cam_idx]: + if not render_rgb[camid] and not render_depth[camid] and not render_seg[camid]: return # Map active camera index to MuJoCo camera ID - mujoco_cam_id = cam_id_map[cam_idx] + mujoco_cam_id = cam_id_map[camid] if wp.static(rc.use_precomputed_rays): ray_dir_local_cam = ray[rayid] else: - img_w = cam_res[cam_idx][0] - img_h = cam_res[cam_idx][1] + img_w = cam_res[camid][0] + img_h = cam_res[camid][1] px = rayid_local % img_w py = rayid_local // img_w ray_dir_local_cam = compute_ray( @@ -742,24 +646,25 @@ def render(m: Model, d: Data, rc: RenderContext): flex_group_root, ray_origin_world, ray_dir_world, + float(MJ_MAXVAL), wp.static(rc.enable_backface_culling), ) - if render_seg[cam_idx] and geom_id != -1: + if render_seg[camid] and geom_id != -1: if geom_id == -2: - seg_out[worldid, seg_adr[cam_idx] + rayid_local] = wp.vec2i(mesh_id, int(ObjType.FLEX)) + seg_out[worldid, seg_adr[camid] + rayid_local] = wp.vec2i(mesh_id, int(ObjType.FLEX)) else: - seg_out[worldid, seg_adr[cam_idx] + rayid_local] = wp.vec2i(geom_id, int(ObjType.GEOM)) + seg_out[worldid, seg_adr[camid] + rayid_local] = wp.vec2i(geom_id, int(ObjType.GEOM)) # Early Out if geom_id == -1: - if wp.static(rc.render_skybox) and render_rgb[cam_idx]: + if wp.static(rc.render_skybox) and render_rgb[camid]: skybox_color = sample_skybox( textures[wp.static(rc.skybox_tex_id)], wp.static(1.0 / float(rc.skybox_face_width)), ray_dir_world, ) - rgb_out[worldid, rgb_adr[cam_idx] + rayid_local] = pack_rgba_to_uint32( + rgb_out[worldid, rgb_adr[camid] + rayid_local] = pack_rgba_to_uint32( skybox_color[0] * 255.0, skybox_color[1] * 255.0, skybox_color[2] * 255.0, @@ -767,14 +672,14 @@ def render(m: Model, d: Data, rc: RenderContext): ) return - if render_depth[cam_idx]: + if render_depth[camid]: # Planar depth: project Euclidean distance onto the camera's optical axis. # In camera-local coordinates, the optical axis is -Z. The Z-component of the # normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ) # between the ray and the optical axis. - depth_out[worldid, depth_adr[cam_idx] + rayid_local] = dist * (-ray_dir_local_cam[2]) + depth_out[worldid, depth_adr[camid] + rayid_local] = dist * (-ray_dir_local_cam[2]) - if not render_rgb[cam_idx]: + if not render_rgb[camid]: return # Shade the pixel @@ -864,7 +769,7 @@ def render(m: Model, d: Data, rc: RenderContext): hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0)) hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0)) - rgb_out[worldid, rgb_adr[cam_idx] + rayid_local] = pack_rgba_to_uint32( + rgb_out[worldid, rgb_adr[camid] + rayid_local] = pack_rgba_to_uint32( hit_color[0] * 255.0, hit_color[1] * 255.0, hit_color[2] * 255.0, @@ -934,4 +839,5 @@ def render(m: Model, d: Data, rc: RenderContext): rc.depth_data, rc.seg_data, ], + block_dim=m.block_dim.render, ) diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py index 61e5b219..97fb1f39 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py @@ -1012,7 +1012,7 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt elif element == State.EQ_ACTIVE: for j in range(neq): state_out[worldid, adr + j] = float(eq_active_in[worldid, j]) - adr += j + adr += neq elif element == State.MOCAP_POS: for j in range(nmocap): pos = mocap_pos_in[worldid, j] @@ -1160,12 +1160,12 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt elif element == State.EQ_ACTIVE: for j in range(neq): eq_active_out[worldid, j] = bool(state_in[worldid, adr + j]) - adr += j + adr += neq elif element == State.MOCAP_POS: for j in range(nmocap): pos = wp.vec3( - state_in[worldid, adr + 1], state_in[worldid, adr + 0], + state_in[worldid, adr + 1], state_in[worldid, adr + 2], ) mocap_pos_out[worldid, j] = pos diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py index 77c6cca2..4ce5a31c 100644 --- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py +++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py @@ -67,6 +67,7 @@ class BlockDim: linesearch_iterative: linesearch iterative block dimension (solver) contact_jac_tiled: contact Jacobian tiled block dimension (solver) qderiv_actuator_dense: qderiv actuator dense block dimension (derivative) + render: render block dimension (render) """ # collision_driver @@ -93,6 +94,8 @@ class BlockDim: contact_jac_tiled: int = 32 # derivative qderiv_actuator_dense: int = 32 + # render + render: int = 64 class BroadphaseType(enum.IntEnum): @@ -2206,6 +2209,8 @@ class RenderContext: mesh-ray rule. When False, the renderer reports inner-surface hits, which is faster but causes a camera placed inside a geom to render that geom's back wall. + geom_ray_types: tuple of GeomType int values present in the scene, used to + statically eliminate unused intersection branches in the ray-cast kernels. """ nrender: int @@ -2260,3 +2265,4 @@ class RenderContext: znear: float total_rays: int enable_backface_culling: bool + geom_ray_types: tuple = () diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 1301aa88..7f648694 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -14,19 +14,17 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" - import dataclasses import functools - import jax -import warp as wp - from mujoco.mjx._src import types -import mujoco.mjx.third_party.mujoco_warp as mjwarp -from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types from mujoco.mjx.warp import ffi from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS from mujoco.mjx.warp.render_context import RenderContextPytree +import mujoco.mjx.third_party.mujoco_warp as mjwarp +from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types +import warp as wp + _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} @@ -50,7 +48,6 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) - @ffi.format_args_for_warp def _refit_bvh_shim( # Model diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index e3d6b595..d3cfd111 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -14,17 +14,14 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" - import dataclasses import functools - import jax -import warp as wp - from mujoco.mjx._src import types +from mujoco.mjx.warp import ffi import mujoco.mjx.third_party.mujoco_warp as mjwarp from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types -from mujoco.mjx.warp import ffi +import warp as wp _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} @@ -48,7 +45,6 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) - @ffi.format_args_for_warp def _collision_shim( # Model diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 8bae7621..5242732d 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -14,17 +14,14 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" - import dataclasses import functools - import jax -import warp as wp - from mujoco.mjx._src import types +from mujoco.mjx.warp import ffi import mujoco.mjx.third_party.mujoco_warp as mjwarp from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types -from mujoco.mjx.warp import ffi +import warp as wp _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} @@ -48,7 +45,6 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) - @ffi.format_args_for_warp def _forward_shim( # Model diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 92311b58..ae82028f 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -14,19 +14,17 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" - import dataclasses import functools - import jax -import warp as wp - from mujoco.mjx._src import types -import mujoco.mjx.third_party.mujoco_warp as mjwarp -from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types from mujoco.mjx.warp import ffi from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS from mujoco.mjx.warp.render_context import RenderContextPytree +import mujoco.mjx.third_party.mujoco_warp as mjwarp +from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types +import warp as wp + _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} @@ -55,6 +53,7 @@ _cb = mjwp_types.Callback( def _render_shim( # Model nworld: int, + block_dim: mjwp_types.BlockDim, cam_fovy: wp.array2d[float], cam_intrinsic: wp.array2d[wp.vec4], cam_projection: wp.array[int], @@ -94,6 +93,7 @@ def _render_shim( _m.callback = _cb _d.efc = _e _d.contact = _c + _m.block_dim = block_dim _m.cam_fovy = cam_fovy _m.cam_intrinsic = cam_intrinsic _m.cam_projection = cam_projection @@ -164,6 +164,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree): ) out = jf( render_ctx.nworld, + m._impl.block_dim, m.cam_fovy, m.cam_intrinsic, m._impl.cam_projection, diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 862b6d07..bf26a123 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -14,17 +14,14 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" - import dataclasses import functools - import jax -import warp as wp - from mujoco.mjx._src import types +from mujoco.mjx.warp import ffi import mujoco.mjx.third_party.mujoco_warp as mjwarp from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types -from mujoco.mjx.warp import ffi +import warp as wp _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} @@ -48,7 +45,6 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) - @ffi.format_args_for_warp def _kinematics_shim( # Model diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 8f952897..cc8f089d 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -15,17 +15,14 @@ """MJX Warp types. DO NOT EDIT. This file is auto-generated. """ - import dataclasses import typing from typing import Tuple - import jax from jax import tree_util from jax.interpreters import batching -import numpy as np - from mujoco.mjx._src import dataclasses as mjx_dataclasses +import numpy as np if typing.TYPE_CHECKING: GraphMode = int @@ -37,7 +34,6 @@ if typing.TYPE_CHECKING: else: try: from warp._src.jax_experimental.ffi import GraphMode - from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types Callback = mjwp_types.Callback @@ -46,7 +42,6 @@ else: Callback = None PyTreeNode = mjx_dataclasses.PyTreeNode - @dataclasses.dataclass(frozen=True) @tree_util.register_pytree_node_class class TileSet: @@ -58,7 +53,6 @@ class TileSet: adr: address of each tile in the set size: size of all the tiles in this set """ - adr: np.ndarray size: int @@ -101,8 +95,8 @@ class BlockDim: linesearch_iterative: linesearch iterative block dimension (solver) contact_jac_tiled: contact Jacobian tiled block dimension (solver) qderiv_actuator_dense: qderiv actuator dense block dimension (derivative) + render: render block dimension (render) """ - actuator_velocity: int cholesky_factorize: int cholesky_factorize_solve: int @@ -114,6 +108,7 @@ class BlockDim: linesearch_iterative: int qderiv_actuator_dense: int ray: int + render: int segmented_sort: int solve_LD_sparse_fused: int update_gradient_JTDAJ_dense: int @@ -133,13 +128,10 @@ class BlockDim: class StatisticWarp(PyTreeNode): """Derived fields from Statistic.""" - meaninertia: jax.Array - class OptionWarp(PyTreeNode): """Derived fields from Option.""" - broadphase: int broadphase_filter: int ccd_iterations: int @@ -154,7 +146,6 @@ class OptionWarp(PyTreeNode): sdf_initpoints: int sdf_iterations: int - class ModelWarp(PyTreeNode): """Derived fields from Model.""" D_colind: np.ndarray @@ -336,7 +327,6 @@ class ModelWarp(PyTreeNode): wrap_site_adr: np.ndarray wrap_site_pair_adr: np.ndarray - class DataWarp(PyTreeNode): """Derived fields from Data.""" M: jax.Array @@ -448,8 +438,6 @@ class DataWarp(PyTreeNode): wrap_obj: jax.Array wrap_xpos: jax.Array shape = property(lambda self: self.cacc.shape) - - DATA_NON_VMAP = { 'contact__dim', 'contact__dist', @@ -478,7 +466,6 @@ DATA_NON_VMAP = { 'nworld', } - def _to_elt(cont, _, d, axis): return DataWarp(**{ f.name: ( @@ -712,6 +699,7 @@ _NDIM = { 'block_dim__linesearch_iterative': 0, 'block_dim__qderiv_actuator_dense': 0, 'block_dim__ray': 0, + 'block_dim__render': 0, 'block_dim__segmented_sort': 0, 'block_dim__solve_LD_sparse_fused': 0, 'block_dim__update_gradient_JTDAJ_dense': 0, @@ -1347,6 +1335,7 @@ _BATCH_DIM = { 'block_dim__linesearch_iterative': False, 'block_dim__qderiv_actuator_dense': False, 'block_dim__ray': False, + 'block_dim__render': False, 'block_dim__segmented_sort': False, 'block_dim__solve_LD_sparse_fused': False, 'block_dim__update_gradient_JTDAJ_dense': False,