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

Merge remote-tracking branch 'origin/main' into tarik/mjx-warp-nested-dataclass-methods

Browse Source 
# Conflicts:
#	mjx/mujoco/mjx/warp/types.py
This commit is contained in:
Tarik Kelestemur
2026-06-03 11:36:53 -04:00
parent c80d8fb35d 0897e3aec1
commit 319b3975a0
68 changed files with 6108 additions and 1631 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/codegen/generate_warp_types.py
+15 -194
View File
@@ -17,23 +17,21 @@
import ast
import dataclasses
import enum
import inspect
import logging
import textwrap
import typing
from typing import Any, Callable, Dict, List, Optional, Set
from absl import app
from absl import flags
from etils import epath
import numpy as np
import warp as wp
import mujoco
from mujoco.mjx.codegen import file
import mujoco.mjx.third_party.mujoco_warp as mjwarp
import numpy as np
import warp as wp
from mujoco.mjx.third_party.warp._src.jax_experimental import ffi
_MJX_WARP_TYPES_OUT_FPATH = flags.DEFINE_string(
'mjx_warp_types_out_path',
'third_party/py/mujoco/mjx/warp/types.py',
@@ -47,7 +45,6 @@ _MJX_TYPES_PATH = flags.DEFINE_string(
)
_DATA_SHAPE_PROPERTY_FIELD = 'cacc'
_DOCSTRING_LINE_LENGTH = 80
_DUMMY_XML = """
<mujoco>
<worldbody>
@@ -64,29 +61,6 @@ _DUMMY_XML = """
"""
def _format_docstring(docstring: str) -> str:
"""Wraps docstring body lines to match checked-in generated files."""
formatted_lines = []
for line in docstring.splitlines():
if not line.strip():
formatted_lines.append(line)
continue
indent = line[: len(line) - len(line.lstrip())]
continuation_indent = indent + (' ' if indent else '')
formatted_lines.append(
textwrap.fill(
line,
width=_DOCSTRING_LINE_LENGTH,
subsequent_indent=continuation_indent,
break_long_words=False,
break_on_hyphens=False,
)
)
return '\n'.join(formatted_lines)
def _to_py_string(value, indent=0):
"""Converts a dictionary/set/tuple/type to a Python code string."""
indent_str = ' ' * indent
@@ -135,10 +109,8 @@ def _get_target_annotation_node(
if annotation == np.ndarray:
return _ast_parse_type('np.ndarray')
if (
isinstance(annotation, wp.array)
or type(annotation).__name__ == '_ArrayAnnotation'
):
if (isinstance(annotation, wp.array) or
type(annotation).__name__ == '_ArrayAnnotation'):
return _ast_parse_type('jax.Array')
if annotation in (int, float, bool):
@@ -194,10 +166,9 @@ def _get_annotations_recursive(
flattened = {}
for key, annotation in annotations.items():
full_key = f'{prefix}{key}'
if (
hasattr(annotation, '__annotations__')
and 'mujoco_warp' in annotation.__module__
):
if hasattr(
annotation, '__annotations__'
) and 'mujoco_warp' in annotation.__module__:
nested = _get_annotations_recursive(
dict(annotation.__annotations__), prefix=f'{full_key}__'
)
@@ -310,12 +281,6 @@ else:
Callback = None
PyTreeNode = mjx_dataclasses.PyTreeNode
def _as_numpy_array(value):
if hasattr(value, 'numpy'):
return value.numpy()
return np.asarray(value)
'''
target_fpath.write_text(header)
@@ -333,68 +298,6 @@ _FLATTEN_UNFLATTEN = """
"""
class _NumpyMethodAdapter(ast.NodeTransformer):
"""Adapts copied Warp array methods to MJX numpy-backed fields."""
def _is_asarray_call(self, node: ast.Call) -> bool:
return (
isinstance(node.func, ast.Attribute)
and isinstance(node.func.value, ast.Name)
and node.func.value.id == 'np'
and node.func.attr == 'asarray'
and len(node.args) == 1
and not node.keywords
)
def _is_as_numpy_array_call(self, node: ast.AST) -> bool:
return (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == '_as_numpy_array'
)
def visit_Call(self, node: ast.Call) -> ast.AST: # pylint: disable=invalid-name
self.generic_visit(node)
if (
isinstance(node.func, ast.Attribute)
and node.func.attr == 'numpy'
and not node.args
and not node.keywords
):
return ast.copy_location(
ast.Call(
func=ast.Name(id='_as_numpy_array', ctx=ast.Load()),
args=[node.func.value],
keywords=[],
),
node,
)
if self._is_asarray_call(node) and self._is_as_numpy_array_call(
node.args[0]
):
return ast.copy_location(node.args[0], node)
return node
def _get_explicit_method_nodes(cls: Any) -> List[ast.FunctionDef]:
"""Returns explicit methods from a source class, adapted for MJX fields."""
source = textwrap.dedent(inspect.getsource(cls))
tree = ast.parse(source)
class_def = next(node for node in tree.body if isinstance(node, ast.ClassDef))
methods: List[ast.FunctionDef] = []
for node in class_def.body:
if not isinstance(node, ast.FunctionDef):
continue
if node.name in {'tree_flatten', 'tree_unflatten'}:
continue
methods.append(
typing.cast(ast.FunctionDef, _NumpyMethodAdapter().visit(node))
)
return methods
def write_nested_dataclass(target_fpath: epath.Path, cls: Any):
new_class_body = _build_new_class_body_ast(
set(cls.__annotations__.keys()),
@@ -402,89 +305,19 @@ def write_nested_dataclass(target_fpath: epath.Path, cls: Any):
dict(cls.__annotations__),
add_docstring=False,
)
new_class_body.extend(_get_explicit_method_nodes(cls))
cls_str = '\n'.join(
textwrap.indent(ast.unparse(node), ' ') for node in new_class_body
)
cls_str = '\n'.join([' ' + ast.unparse(node) for node in new_class_body])
cls_str = cls_str.replace('jax.Array', 'np.ndarray')
docstring = _format_docstring(cls.__doc__ or '')
with target_fpath.open('a') as f:
f.write(f'''
@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class {cls.__name__}:
"""{docstring}"""
"""{cls.__doc__}"""
{cls_str}
{_FLATTEN_UNFLATTEN}
''')
def _get_class_name(line: str) -> str | None:
"""Returns the class name from a generated class definition line."""
if not line.startswith('class '):
return None
return line.removeprefix('class ').split('(', maxsplit=1)[0]
def _get_compact_derived_docstring_classes(
target_fpath: epath.Path,
) -> Set[str]:
"""Finds generated classes with no blank line after their docstring."""
compact_classes = set()
if not target_fpath.exists():
return compact_classes
lines = target_fpath.read_text().splitlines()
for i, line in enumerate(lines[:-2]):
cls_name = _get_class_name(line)
if cls_name is None:
continue
if (
'(PyTreeNode)' in line
and lines[i + 1].startswith(' """Derived fields from ')
and lines[i + 1].rstrip().endswith('."""')
and lines[i + 2].startswith(' ')
and ':' in lines[i + 2]
):
compact_classes.add(cls_name)
return compact_classes
def _restore_compact_derived_docstring_spacing(
target_fpath: epath.Path,
compact_classes: Set[str],
):
"""Restores compact generated PyTreeNode field class spacing."""
if not compact_classes:
return
lines = target_fpath.read_text().splitlines(keepends=True)
result = []
i = 0
current_class = None
while i < len(lines):
cls_name = _get_class_name(lines[i])
if cls_name is not None:
current_class = cls_name
result.append(lines[i])
if (
current_class in compact_classes
and lines[i].startswith(' """Derived fields from ')
and lines[i].rstrip().endswith('."""')
and i + 2 < len(lines)
and lines[i + 1].strip() == ''
and lines[i + 2].startswith(' ')
and ':' in lines[i + 2]
):
i += 2
continue
i += 1
target_fpath.write_text(''.join(result))
def _get_meta_fields(cls_name: str) -> Set[str]:
"""Returns the set of fields that should be meta-fields in the pytree."""
m = mujoco.MjModel.from_xml_string(_DUMMY_XML)
@@ -597,10 +430,8 @@ def _get_fields_with_cond(
if f.type in (int, float, bool) and add_static:
s.add(prefix + f.name)
continue
if not (
isinstance(f.type, wp.array)
or type(f.type).__name__ == '_ArrayAnnotation'
):
if not (isinstance(f.type, wp.array) or
type(f.type).__name__ == '_ArrayAnnotation'):
continue
if cond_fn(attr):
s.add(prefix + f.name)
@@ -639,10 +470,8 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
def _is_ffi_compatible(wp_type: Any) -> bool:
"""Returns True if the type is an array, scalar, or variadic tuple."""
if (
isinstance(wp_type, wp.array)
or type(wp_type).__name__ == '_ArrayAnnotation'
):
if (isinstance(wp_type, wp.array) or
type(wp_type).__name__ == '_ArrayAnnotation'):
return True
if wp_type in wp._src.types.value_types:
return True
@@ -719,9 +548,6 @@ def main(argv):
base_path = file.get_base_path()
target_fpath = base_path / _MJX_WARP_TYPES_OUT_FPATH.value
mjx_types_fpath = base_path / _MJX_TYPES_PATH.value
compact_derived_docstring_classes = _get_compact_derived_docstring_classes(
target_fpath
)
write_header(target_fpath)
# TODO(btaba): consider automated grabbing of nested dataclasses from mjwarp.
@@ -730,9 +556,7 @@ def main(argv):
write_core_cls('Statistic', target_fpath, mjx_types_fpath, set_diff=False)
write_core_cls(
'Option',
target_fpath,
mjx_types_fpath,
'Option', target_fpath, mjx_types_fpath,
extra_annotations={'graph_mode': ffi.GraphMode},
)
write_core_cls('Model', target_fpath, mjx_types_fpath)
@@ -743,9 +567,6 @@ def main(argv):
file.write_license(target_fpath)
file.format_file(target_fpath)
_restore_compact_derived_docstring_spacing(
target_fpath, compact_derived_docstring_classes
)
if __name__ == '__main__':
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
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)
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])
# 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,
)
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:
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
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)
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 = ()
mjx/mujoco/mjx/warp/bvh.py
+4 -7
View File
@@ -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
mjx/mujoco/mjx/warp/collision_driver.py
+2 -6
View File
@@ -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
mjx/mujoco/mjx/warp/forward.py
+2 -6
View File
@@ -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
mjx/mujoco/mjx/warp/render.py
+7 -6
View File
@@ -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,
mjx/mujoco/mjx/warp/smooth.py
+2 -6
View File
@@ -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
mjx/mujoco/mjx/warp/types.py
+9 -24
View File
@@ -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,13 +42,6 @@ else:
Callback = None
PyTreeNode = mjx_dataclasses.PyTreeNode
def _as_numpy_array(value):
if hasattr(value, 'numpy'):
return value.numpy()
return np.asarray(value)
@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class TileSet:
@@ -64,19 +53,20 @@ class TileSet:
adr: address of each tile in the set
size: size of all the tiles in this set
"""
adr: np.ndarray
size: int
# Manually kept in this generated shim until TileSet method generation is
# needed more broadly. Keep this in sync with mujoco_warp._src.types.TileSet.
def __eq__(self, other) -> bool:
if self.__class__ is not other.__class__:
return NotImplemented
return self.size == other.size and np.array_equal(
_as_numpy_array(self.adr), _as_numpy_array(other.adr)
np.asarray(self.adr), np.asarray(other.adr)
)
def __hash__(self) -> int:
adr = _as_numpy_array(self.adr)
adr = np.asarray(self.adr)
return hash((self.size, adr.dtype.str, adr.shape, adr.tobytes()))
def tree_flatten(self):
@@ -118,8 +108,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
@@ -131,6 +121,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
@@ -150,13 +141,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
@@ -171,7 +159,6 @@ class OptionWarp(PyTreeNode):
sdf_initpoints: int
sdf_iterations: int
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
D_colind: np.ndarray
@@ -353,7 +340,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
@@ -465,8 +451,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',
@@ -495,7 +479,6 @@ DATA_NON_VMAP = {
'nworld',
}
def _to_elt(cont, _, d, axis):
return DataWarp(**{
f.name: (
@@ -729,6 +712,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,
@@ -1364,6 +1348,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,