chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 925939529
Change-Id: Idb92c2d50d5be9df7b6e91113b576e9c706847ca
This commit is contained in:
Taylor Howell
2026-06-03 04:40:22 -07:00
committed by Copybara-Service
parent 4cfdf5f34b
commit 5d231c64cd
10 changed files with 379 additions and 484 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+10 -2
View File
@@ -3058,11 +3058,11 @@ def create_render_context(
# Locate skybox texture # Locate skybox texture
skybox_tex_ids = np.nonzero(mjm.tex_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX)[0] if mjm.ntex else np.array([], dtype=int) skybox_tex_ids = np.nonzero(mjm.tex_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX)[0] if mjm.ntex else np.array([], dtype=int)
if render_skybox: if render_skybox and skybox_tex_ids.size > 0:
assert skybox_tex_ids.size > 0, "render_skybox=True but the model has no texture with type mjTEXTURE_SKYBOX"
skybox_tex_id = int(skybox_tex_ids[0]) skybox_tex_id = int(skybox_tex_ids[0])
skybox_face_width = int(mjm.tex_width[skybox_tex_id]) skybox_face_width = int(mjm.tex_width[skybox_tex_id])
else: else:
render_skybox = False
skybox_tex_id = -1 skybox_tex_id = -1
skybox_face_width = 1 skybox_face_width = 1
@@ -3157,6 +3157,13 @@ def create_render_context(
bvh_ngeom = len(geom_enabled_idx) 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( rc = types.RenderContext(
nrender=ncam, nrender=ncam,
cam_res=cam_res_arr, cam_res=cam_res_arr,
@@ -3210,6 +3217,7 @@ def create_render_context(
znear=znear, znear=znear,
total_rays=int(total), total_rays=int(total),
enable_backface_culling=enable_backface_culling, enable_backface_culling=enable_backface_culling,
geom_ray_types=geom_ray_types,
) )
bvh.build_scene_bvh(mjm, mjd, rc, nworld) bvh.build_scene_bvh(mjm, mjd, rc, nworld)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+338 -432
View File
@@ -151,449 +151,346 @@ def sample_skybox(
return wp.vec3(color[0], color[1], color[2]) return wp.vec3(color[0], color[1], color[2])
# TODO: Investigate combining cast_ray and cast_ray_first_hit def _make_cast_ray(geom_ray_types: Tuple[int], first_hit: bool = False) -> wp.Function:
@wp.func """Build a ray-cast func specialized to the geom types present in the scene.
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)
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root) geom_ray_types is the set of GeomType int values that actually occur, so the
bounds_nr = int(0) per-type intersection branches for absent types are eliminated at compile time
ngeom = bvh_ngeom + flex_bvh_ngeom via wp.static, avoiding the register pressure of unreachable code paths.
while wp.bvh_query_next(query, bounds_nr, dist): first_hit selects the variant (also resolved at compile time via wp.static):
gi_global = bounds_nr - False: full closest-hit cast. Returns the closest hit's full surface data.
local_id = gi_global - (worldid * ngeom) - 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) @wp.func
hit_mesh_id = int(-1) def cast_ray(
u = float(0.0) # Model:
v = float(0.0) geom_type: wp.array[int],
f = int(-1) geom_dataid: wp.array2d[int],
n = wp.vec3(0.0, 0.0, 0.0) geom_size: wp.array2d[wp.vec3],
hit_geom_id = int(-1) 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: query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
gi = enabled_geom_ids[local_id] bounds_nr = int(0)
gtype = geom_type[gi] ngeom = bvh_ngeom + flex_bvh_ngeom
else:
gi = local_id - bvh_ngeom
gtype = GeomType.FLEX
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 d = float(-1.0)
if gtype == GeomType.PLANE: hit_mesh_id = int(-1)
d, n = ray_plane( u = float(0.0)
geom_xpos_in[worldid, gi], v = float(0.0)
geom_xmat_in[worldid, gi], f = int(-1)
geom_size[worldid % geom_size.shape[0], gi], n = wp.vec3(0.0, 0.0, 0.0)
ray_origin_world, hit_geom_id = int(-1)
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]
if edge_id >= 0: if local_id < bvh_ngeom:
edge = flex_edge[edge_id] gi = enabled_geom_ids[local_id]
vert_adr = flex_vertadr[flexid] gtype = geom_type[gi]
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
else: else:
flex_gr = flex_group_root[worldid, flexid] gi = local_id - bvh_ngeom
d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist) gtype = GeomType.FLEX
if d >= 0.0:
hit_mesh_id = flexid
# Backface cull: drop exit-face hits when the ray origin is inside the geom, hit_geom_id = gi
# 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
if d >= 0.0 and d < dist: if wp.static(int(GeomType.PLANE) in geom_ray_types):
dist = d if gtype == GeomType.PLANE:
normal = n d, n = ray_plane(
geom_id = hit_geom_id geom_xpos_in[worldid, gi],
bary_u = u geom_xmat_in[worldid, gi],
bary_v = v geom_size[worldid % geom_size.shape[0], gi],
face_idx = f ray_origin_world,
geom_mesh_id = hit_mesh_id 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 d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
def cast_ray_first_hit( hit_mesh_id = flexid
# Model: else:
geom_type: wp.array[int], if wp.static(first_hit):
geom_dataid: wp.array2d[int], hit = ray_flex_with_bvh_anyhit(
geom_size: wp.array2d[wp.vec3], flex_bvh_id,
flex_vertadr: wp.array[int], flexid,
flex_edge: wp.array[wp.vec2i], flex_group_root[worldid, flexid],
flex_radius: wp.array[float], ray_origin_world,
# Data in: ray_dir_world,
geom_xpos_in: wp.array2d[wp.vec3], dist,
geom_xmat_in: wp.array2d[wp.mat33], )
flexvert_xpos_in: wp.array2d[wp.vec3], d = 0.0 if hit else -1.0
# In: else:
bvh_id: wp.uint64, flex_gr = flex_group_root[worldid, flexid]
group_root: int, d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist)
worldid: int, if d >= 0.0:
bvh_ngeom: int, hit_mesh_id = flexid
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
while wp.bvh_query_next(query, bounds_nr, max_dist): # Backface cull: drop exit-face hits when the ray origin is inside the geom,
gi_global = bounds_nr # matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule. Strict `> 0` keeps
local_id = gi_global - (worldid * ngeom) # 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) if wp.static(first_hit):
n = wp.vec3(0.0, 0.0, 0.0) # Any-hit: return as soon as anything is in range; surface data is unused.
if d >= 0.0 and d < dist:
if local_id < bvh_ngeom: return hit_geom_id, d, n, u, v, f, hit_mesh_id
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)
else: else:
hit = ray_flex_with_bvh_anyhit( if d >= 0.0 and d < dist:
flex_bvh_id, dist = d
flexid, normal = n
flex_group_root[worldid, flexid], geom_id = hit_geom_id
ray_origin_world, bary_u = u
ray_dir_world, bary_v = v
max_dist, face_idx = f
) geom_mesh_id = hit_mesh_id
d = 0.0 if hit else -1.0
# Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id
# 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
if d >= 0.0 and d < max_dist: return cast_ray
return True
return False
@wp.func def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function:
def compute_lighting( """Build specialized 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)
# TODO: We should probably only be looping over active lights @wp.func
# in the first place with a static loop of enabled light idx? def compute_lighting(
if not lightactive: # Model:
return light_contribution 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) # TODO: We should probably only be looping over active lights
dist_to_light = float(MJ_MAXVAL) # in the first place with a static loop of enabled light idx?
attenuation = float(1.0) if not lightactive:
return light_contribution
if lighttype == 1: # directional light L = wp.vec3(0.0, 0.0, 0.0)
L = wp.normalize(-lightdir) dist_to_light = float(MJ_MAXVAL)
else: attenuation = float(1.0)
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
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 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( ndotl = wp.max(0.0, wp.dot(normal, L))
geom_type, if ndotl == 0.0:
geom_dataid, return light_contribution
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_hit: visible = float(1.0)
visible = 0.3
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 @event_scope
@@ -611,6 +508,13 @@ def render(m: Model, d: Data, rc: RenderContext):
rc.depth_data.fill_(0.0) rc.depth_data.fill_(0.0)
rc.seg_data.fill_(wp.vec2i(-1, -1)) 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) @wp.kernel(module="unique", enable_backward=False)
def _render_megakernel( def _render_megakernel(
# Model: # Model:
@@ -676,31 +580,31 @@ def render(m: Model, d: Data, rc: RenderContext):
): ):
worldid, rayid = wp.tid() worldid, rayid = wp.tid()
# Map global rayid -> (cam_idx, rayid_local) using cumulative sizes # Map global rayid -> (camid, rayid_local) using cumulative sizes
cam_idx = int(-1) camid = int(-1)
rayid_local = int(-1) rayid_local = int(-1)
accum = int(0) accum = int(0)
for i in range(nrender): for i in range(nrender):
num_i = cam_res[i][0] * cam_res[i][1] num_i = cam_res[i][0] * cam_res[i][1]
if rayid < accum + num_i: if rayid < accum + num_i:
cam_idx = i camid = i
rayid_local = rayid - accum rayid_local = rayid - accum
break break
accum += num_i accum += num_i
if cam_idx == -1 or rayid_local < 0: if camid == -1 or rayid_local < 0:
return 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 return
# Map active camera index to MuJoCo camera ID # 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): if wp.static(rc.use_precomputed_rays):
ray_dir_local_cam = ray[rayid] ray_dir_local_cam = ray[rayid]
else: else:
img_w = cam_res[cam_idx][0] img_w = cam_res[camid][0]
img_h = cam_res[cam_idx][1] img_h = cam_res[camid][1]
px = rayid_local % img_w px = rayid_local % img_w
py = rayid_local // img_w py = rayid_local // img_w
ray_dir_local_cam = compute_ray( ray_dir_local_cam = compute_ray(
@@ -742,24 +646,25 @@ def render(m: Model, d: Data, rc: RenderContext):
flex_group_root, flex_group_root,
ray_origin_world, ray_origin_world,
ray_dir_world, ray_dir_world,
float(MJ_MAXVAL),
wp.static(rc.enable_backface_culling), 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: 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: 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 # Early Out
if geom_id == -1: 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( skybox_color = sample_skybox(
textures[wp.static(rc.skybox_tex_id)], textures[wp.static(rc.skybox_tex_id)],
wp.static(1.0 / float(rc.skybox_face_width)), wp.static(1.0 / float(rc.skybox_face_width)),
ray_dir_world, 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[0] * 255.0,
skybox_color[1] * 255.0, skybox_color[1] * 255.0,
skybox_color[2] * 255.0, skybox_color[2] * 255.0,
@@ -767,14 +672,14 @@ def render(m: Model, d: Data, rc: RenderContext):
) )
return return
if render_depth[cam_idx]: if render_depth[camid]:
# Planar depth: project Euclidean distance onto the camera's optical axis. # 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 # 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(θ) # normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ)
# between the ray and the optical axis. # 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 return
# Shade the pixel # 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.min(result, wp.vec3(1.0, 1.0, 1.0))
hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.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[0] * 255.0,
hit_color[1] * 255.0, hit_color[1] * 255.0,
hit_color[2] * 255.0, hit_color[2] * 255.0,
@@ -934,4 +839,5 @@ def render(m: Model, d: Data, rc: RenderContext):
rc.depth_data, rc.depth_data,
rc.seg_data, rc.seg_data,
], ],
block_dim=m.block_dim.render,
) )
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+3 -3
View File
@@ -1012,7 +1012,7 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
elif element == State.EQ_ACTIVE: elif element == State.EQ_ACTIVE:
for j in range(neq): for j in range(neq):
state_out[worldid, adr + j] = float(eq_active_in[worldid, j]) state_out[worldid, adr + j] = float(eq_active_in[worldid, j])
adr += j adr += neq
elif element == State.MOCAP_POS: elif element == State.MOCAP_POS:
for j in range(nmocap): for j in range(nmocap):
pos = mocap_pos_in[worldid, j] 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: elif element == State.EQ_ACTIVE:
for j in range(neq): for j in range(neq):
eq_active_out[worldid, j] = bool(state_in[worldid, adr + j]) eq_active_out[worldid, j] = bool(state_in[worldid, adr + j])
adr += j adr += neq
elif element == State.MOCAP_POS: elif element == State.MOCAP_POS:
for j in range(nmocap): for j in range(nmocap):
pos = wp.vec3( pos = wp.vec3(
state_in[worldid, adr + 1],
state_in[worldid, adr + 0], state_in[worldid, adr + 0],
state_in[worldid, adr + 1],
state_in[worldid, adr + 2], state_in[worldid, adr + 2],
) )
mocap_pos_out[worldid, j] = pos mocap_pos_out[worldid, j] = pos
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+6
View File
@@ -67,6 +67,7 @@ class BlockDim:
linesearch_iterative: linesearch iterative block dimension (solver) linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver) contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative) qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
render: render block dimension (render)
""" """
# collision_driver # collision_driver
@@ -93,6 +94,8 @@ class BlockDim:
contact_jac_tiled: int = 32 contact_jac_tiled: int = 32
# derivative # derivative
qderiv_actuator_dense: int = 32 qderiv_actuator_dense: int = 32
# render
render: int = 64
class BroadphaseType(enum.IntEnum): class BroadphaseType(enum.IntEnum):
@@ -2206,6 +2209,8 @@ class RenderContext:
mesh-ray rule. When False, the renderer reports inner-surface hits, which 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 is faster but causes a camera placed inside a geom to render that geom's
back wall. 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 nrender: int
@@ -2260,3 +2265,4 @@ class RenderContext:
znear: float znear: float
total_rays: int total_rays: int
enable_backface_culling: bool enable_backface_culling: bool
geom_ray_types: tuple = ()
mjx/mujoco/mjx/warp/bvh.py
+4 -7
View File
@@ -14,19 +14,17 @@
# ============================================================================== # ==============================================================================
"""DO NOT EDIT. This file is auto-generated.""" """DO NOT EDIT. This file is auto-generated."""
import dataclasses import dataclasses
import functools import functools
import jax import jax
import warp as wp
from mujoco.mjx._src import types 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 import ffi
from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.render_context import RenderContextPytree 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( _m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} **{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} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _refit_bvh_shim( def _refit_bvh_shim(
# Model # Model
mjx/mujoco/mjx/warp/collision_driver.py
+2 -6
View File
@@ -14,17 +14,14 @@
# ============================================================================== # ==============================================================================
"""DO NOT EDIT. This file is auto-generated.""" """DO NOT EDIT. This file is auto-generated."""
import dataclasses import dataclasses
import functools import functools
import jax import jax
import warp as wp
from mujoco.mjx._src import types from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp 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.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi import warp as wp
_m = mjwarp.Model( _m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} **{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} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _collision_shim( def _collision_shim(
# Model # Model
mjx/mujoco/mjx/warp/forward.py
+2 -6
View File
@@ -14,17 +14,14 @@
# ============================================================================== # ==============================================================================
"""DO NOT EDIT. This file is auto-generated.""" """DO NOT EDIT. This file is auto-generated."""
import dataclasses import dataclasses
import functools import functools
import jax import jax
import warp as wp
from mujoco.mjx._src import types from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp 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.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi import warp as wp
_m = mjwarp.Model( _m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} **{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} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _forward_shim( def _forward_shim(
# Model # Model
mjx/mujoco/mjx/warp/render.py
+7 -6
View File
@@ -14,19 +14,17 @@
# ============================================================================== # ==============================================================================
"""DO NOT EDIT. This file is auto-generated.""" """DO NOT EDIT. This file is auto-generated."""
import dataclasses import dataclasses
import functools import functools
import jax import jax
import warp as wp
from mujoco.mjx._src import types 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 import ffi
from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.render_context import RenderContextPytree 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( _m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -55,6 +53,7 @@ _cb = mjwp_types.Callback(
def _render_shim( def _render_shim(
# Model # Model
nworld: int, nworld: int,
block_dim: mjwp_types.BlockDim,
cam_fovy: wp.array2d[float], cam_fovy: wp.array2d[float],
cam_intrinsic: wp.array2d[wp.vec4], cam_intrinsic: wp.array2d[wp.vec4],
cam_projection: wp.array[int], cam_projection: wp.array[int],
@@ -94,6 +93,7 @@ def _render_shim(
_m.callback = _cb _m.callback = _cb
_d.efc = _e _d.efc = _e
_d.contact = _c _d.contact = _c
_m.block_dim = block_dim
_m.cam_fovy = cam_fovy _m.cam_fovy = cam_fovy
_m.cam_intrinsic = cam_intrinsic _m.cam_intrinsic = cam_intrinsic
_m.cam_projection = cam_projection _m.cam_projection = cam_projection
@@ -164,6 +164,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
) )
out = jf( out = jf(
render_ctx.nworld, render_ctx.nworld,
m._impl.block_dim,
m.cam_fovy, m.cam_fovy,
m.cam_intrinsic, m.cam_intrinsic,
m._impl.cam_projection, m._impl.cam_projection,
mjx/mujoco/mjx/warp/smooth.py
+2 -6
View File
@@ -14,17 +14,14 @@
# ============================================================================== # ==============================================================================
"""DO NOT EDIT. This file is auto-generated.""" """DO NOT EDIT. This file is auto-generated."""
import dataclasses import dataclasses
import functools import functools
import jax import jax
import warp as wp
from mujoco.mjx._src import types from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp 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.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi import warp as wp
_m = mjwarp.Model( _m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} **{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} **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _kinematics_shim( def _kinematics_shim(
# Model # Model
mjx/mujoco/mjx/warp/types.py
+5 -16
View File
@@ -15,17 +15,14 @@
"""MJX Warp types. """MJX Warp types.
DO NOT EDIT. This file is auto-generated. DO NOT EDIT. This file is auto-generated.
""" """
import dataclasses import dataclasses
import typing import typing
from typing import Tuple from typing import Tuple
import jax import jax
from jax import tree_util from jax import tree_util
from jax.interpreters import batching from jax.interpreters import batching
import numpy as np
from mujoco.mjx._src import dataclasses as mjx_dataclasses from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
if typing.TYPE_CHECKING: if typing.TYPE_CHECKING:
GraphMode = int GraphMode = int
@@ -37,7 +34,6 @@ if typing.TYPE_CHECKING:
else: else:
try: try:
from warp._src.jax_experimental.ffi import GraphMode from warp._src.jax_experimental.ffi import GraphMode
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
Callback = mjwp_types.Callback Callback = mjwp_types.Callback
@@ -46,7 +42,6 @@ else:
Callback = None Callback = None
PyTreeNode = mjx_dataclasses.PyTreeNode PyTreeNode = mjx_dataclasses.PyTreeNode
@dataclasses.dataclass(frozen=True) @dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class @tree_util.register_pytree_node_class
class TileSet: class TileSet:
@@ -58,7 +53,6 @@ class TileSet:
adr: address of each tile in the set adr: address of each tile in the set
size: size of all the tiles in this set size: size of all the tiles in this set
""" """
adr: np.ndarray adr: np.ndarray
size: int size: int
@@ -101,8 +95,8 @@ class BlockDim:
linesearch_iterative: linesearch iterative block dimension (solver) linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver) contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative) qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
render: render block dimension (render)
""" """
actuator_velocity: int actuator_velocity: int
cholesky_factorize: int cholesky_factorize: int
cholesky_factorize_solve: int cholesky_factorize_solve: int
@@ -114,6 +108,7 @@ class BlockDim:
linesearch_iterative: int linesearch_iterative: int
qderiv_actuator_dense: int qderiv_actuator_dense: int
ray: int ray: int
render: int
segmented_sort: int segmented_sort: int
solve_LD_sparse_fused: int solve_LD_sparse_fused: int
update_gradient_JTDAJ_dense: int update_gradient_JTDAJ_dense: int
@@ -133,13 +128,10 @@ class BlockDim:
class StatisticWarp(PyTreeNode): class StatisticWarp(PyTreeNode):
"""Derived fields from Statistic.""" """Derived fields from Statistic."""
meaninertia: jax.Array meaninertia: jax.Array
class OptionWarp(PyTreeNode): class OptionWarp(PyTreeNode):
"""Derived fields from Option.""" """Derived fields from Option."""
broadphase: int broadphase: int
broadphase_filter: int broadphase_filter: int
ccd_iterations: int ccd_iterations: int
@@ -154,7 +146,6 @@ class OptionWarp(PyTreeNode):
sdf_initpoints: int sdf_initpoints: int
sdf_iterations: int sdf_iterations: int
class ModelWarp(PyTreeNode): class ModelWarp(PyTreeNode):
"""Derived fields from Model.""" """Derived fields from Model."""
D_colind: np.ndarray D_colind: np.ndarray
@@ -336,7 +327,6 @@ class ModelWarp(PyTreeNode):
wrap_site_adr: np.ndarray wrap_site_adr: np.ndarray
wrap_site_pair_adr: np.ndarray wrap_site_pair_adr: np.ndarray
class DataWarp(PyTreeNode): class DataWarp(PyTreeNode):
"""Derived fields from Data.""" """Derived fields from Data."""
M: jax.Array M: jax.Array
@@ -448,8 +438,6 @@ class DataWarp(PyTreeNode):
wrap_obj: jax.Array wrap_obj: jax.Array
wrap_xpos: jax.Array wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape) shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = { DATA_NON_VMAP = {
'contact__dim', 'contact__dim',
'contact__dist', 'contact__dist',
@@ -478,7 +466,6 @@ DATA_NON_VMAP = {
'nworld', 'nworld',
} }
def _to_elt(cont, _, d, axis): def _to_elt(cont, _, d, axis):
return DataWarp(**{ return DataWarp(**{
f.name: ( f.name: (
@@ -712,6 +699,7 @@ _NDIM = {
'block_dim__linesearch_iterative': 0, 'block_dim__linesearch_iterative': 0,
'block_dim__qderiv_actuator_dense': 0, 'block_dim__qderiv_actuator_dense': 0,
'block_dim__ray': 0, 'block_dim__ray': 0,
'block_dim__render': 0,
'block_dim__segmented_sort': 0, 'block_dim__segmented_sort': 0,
'block_dim__solve_LD_sparse_fused': 0, 'block_dim__solve_LD_sparse_fused': 0,
'block_dim__update_gradient_JTDAJ_dense': 0, 'block_dim__update_gradient_JTDAJ_dense': 0,
@@ -1347,6 +1335,7 @@ _BATCH_DIM = {
'block_dim__linesearch_iterative': False, 'block_dim__linesearch_iterative': False,
'block_dim__qderiv_actuator_dense': False, 'block_dim__qderiv_actuator_dense': False,
'block_dim__ray': False, 'block_dim__ray': False,
'block_dim__render': False,
'block_dim__segmented_sort': False, 'block_dim__segmented_sort': False,
'block_dim__solve_LD_sparse_fused': False, 'block_dim__solve_LD_sparse_fused': False,
'block_dim__update_gradient_JTDAJ_dense': False, 'block_dim__update_gradient_JTDAJ_dense': False,