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: 960411406
Change-Id: I687b06586881d1915bab1bf8ebab6473e6ef29db
This commit is contained in:
Taylor Howell
2026-08-06 11:23:39 -07:00
committed by Copybara-Service
parent b3f8f91ec5
commit a1d772c9ad
40 changed files with 4793 additions and 3146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/codegen/generate_warp_shim.py
+15 -18
View File
@@ -26,6 +26,7 @@ from etils import epath
from mujoco.mjx._src import types as mjx_types
from mujoco.mjx.codegen import file
from mujoco.mjx.codegen import trace
from mujoco.mjx.warp import types as mjxw_types
import jax
from mujoco.mjx.third_party import mujoco_warp # pylint: disable=unused-import
@@ -79,17 +80,16 @@ def _clean_type(type_: str):
'BlockDim',
'vec_pluginattr',
)
m = re.match(r'array\((.*)\)', type_)
if m: # match custom mujoco_warp array annotation types
args_str = m.group(1)
args = [a.strip() for a in args_str.split(',')]
# Match and convert custom mujoco_warp array annotations (e.g. within tuple[...]).
def _tuple_array(m):
args = [a.strip() for a in m.group(1).split(',')]
ndim, dtype = len(args) - 1, args[-1]
dims = {1: '', 2: '2d', 3: '3d', 4: '4d'}
if ndim not in dims:
raise ValueError(f'Unsupported array ndim: {ndim} for type: {dtype}')
return f'wp.array{dims[ndim]}[{dtype}]'
type_ = f'wp.array{dims[ndim]}[{dtype}]'
type_ = re.sub(r'array\((.*?)\)', _tuple_array, type_)
for t in types_to_prefix:
type_ = re.sub(rf'\b{t}\b', f'mjwp_types.{t}', type_)
@@ -123,31 +123,28 @@ def _get_stage_fields(
return False
return cls.__annotations__.get(field) is jax.Array
ModelWarp = getattr(mjx_types, 'ModelWarp', None)
OptionWarp = getattr(mjx_types, 'OptionWarp', None)
ModelWarp = getattr(mjxw_types, 'ModelWarp', None)
OptionWarp = getattr(mjxw_types, 'OptionWarp', None)
DataWarp = getattr(mjxw_types, 'DataWarp', None)
# stage_in: Model/ModelWarp jax.Array input fields
for field in field_usage.model_fields:
if is_jax_array(mjx_types.Model, field):
stage_in.append(field)
elif is_jax_array(ModelWarp, field):
if is_jax_array(mjx_types.Model, field) or is_jax_array(ModelWarp, field):
stage_in.append(field)
# stage_in: Option/OptionWarp jax.Array input fields
elif field.startswith('opt__'):
sub_field = field.split('opt__')[-1]
if is_jax_array(mjx_types.Option, sub_field):
stage_in.append(field)
elif is_jax_array(OptionWarp, sub_field):
if is_jax_array(mjx_types.Option, sub_field) or is_jax_array(OptionWarp, sub_field):
stage_in.append(field)
# stage_in: Data jax.Array input fields
# stage_in: Data/DataWarp jax.Array input fields
for field in field_usage.data_fields:
if is_jax_array(mjx_types.Data, field):
if is_jax_array(mjx_types.Data, field) or is_jax_array(DataWarp, field):
stage_in.append(field)
# stage_out: Data jax.Array output fields
# stage_out: Data/DataWarp jax.Array output fields
for field in field_usage.data_out_fields:
if is_jax_array(mjx_types.Data, field):
if is_jax_array(mjx_types.Data, field) or is_jax_array(DataWarp, field):
stage_out.append(field)
return sorted(stage_in), sorted(stage_out)
mjx/mujoco/mjx/codegen/generate_warp_types.py
+69 -121
View File
@@ -19,19 +19,17 @@ import dataclasses
import enum
import logging
import typing
from typing import Any, Callable, Dict, List, Optional, Set
from typing import Any, Dict, List, Optional, Set
from absl import app
from absl import flags
from etils import epath
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 warp._src.jax.ffi import FfiArg
from warp import JaxCallableGraphMode as GraphMode
from warp._src.jax.ffi import FfiArg
_MJX_WARP_TYPES_OUT_FPATH = flags.DEFINE_string(
'mjx_warp_types_out_path',
@@ -46,20 +44,33 @@ _MJX_TYPES_PATH = flags.DEFINE_string(
)
_DATA_SHAPE_PROPERTY_FIELD = 'cacc'
_DUMMY_XML = """
<mujoco>
<worldbody>
<body>
<joint type="free"/>
<geom pos="0 0 0" size="0.2" type="sphere"/>
</body>
<body >
<joint type="free"/>
<geom pos="0 0.3 0" size="0.11" type="sphere"/>
</body>
</worldbody>
</mujoco>
"""
_CLS_MAP = {
'Model': mjwarp.Model,
'Data': mjwarp.Data,
'Option': mjwarp.Option,
'Statistic': mjwarp.Statistic,
}
def _is_array_spec(typ) -> bool:
"""Check if a type annotation is a Warp array spec."""
return isinstance(typ, wp.array) or type(typ).__name__ == '_ArrayAnnotation'
def _is_batched_field(field_type) -> Optional[bool]:
"""Returns whether a field annotation is batched.
Mirrors mujoco_warp._src.io._mark_batched: inspects spec_shape[0].
Returns True if batched, False if array but not batched, None if not array.
"""
if not _is_array_spec(field_type):
return None
spec_shape = getattr(field_type, 'shape', ())
if not spec_shape:
return False
return spec_shape[0] in ('*', 'nworld')
def _to_py_string(value, indent=0):
@@ -110,8 +121,7 @@ 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 _is_array_spec(annotation):
return _ast_parse_type('jax.Array')
if annotation in (int, float, bool):
@@ -137,10 +147,7 @@ def _get_target_annotation_node(
type_ = typing.get_args(annotation)[0].__name__
return _ast_parse_type(f'Tuple[{type_}, ...]')
if is_tuple and (
isinstance(typing.get_args(annotation)[0], wp.array)
or type(typing.get_args(annotation)[0]).__name__ == '_ArrayAnnotation'
):
if is_tuple and _is_array_spec(typing.get_args(annotation)[0]):
return _ast_parse_type('Tuple[np.ndarray, ...]')
if is_tuple and dataclasses.is_dataclass(typing.get_args(annotation)[0]):
@@ -349,6 +356,11 @@ def write_nested_dataclass(target_fpath: epath.Path, cls: Any):
)
cls_str = '\n'.join([' ' + ast.unparse(node) for node in new_class_body])
cls_str = cls_str.replace('jax.Array', 'np.ndarray')
if cls.__name__ == 'TileSet':
cls_str = cls_str.replace(
'elemid: np.ndarray',
'elemid: np.ndarray = dataclasses.field(default_factory=lambda: np.array([], dtype=np.int32))'
)
manual_methods = _NESTED_DATACLASS_MANUAL_METHODS.get(cls.__name__, '')
with target_fpath.open('a') as f:
f.write(f'''
@@ -362,33 +374,25 @@ class {cls.__name__}:
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)
d = mujoco.MjData(m)
mujoco.mj_step(m, d)
"""Returns the set of fields that should be meta-fields in the pytree.
Meta-fields are array fields that are not batched.
"""
cls = _CLS_MAP[cls_name]
annotations = _get_annotations_recursive(dict(cls.__annotations__))
# Non-batched array fields (same check as mujoco_warp._src.io._mark_batched).
meta_fields = {
name
for name, type_ in annotations.items()
if _is_batched_field(type_) is False
}
# Data meta-fields exclude non-vmap fields (scalars and non-nworld arrays).
if cls_name == 'Data':
with wp.ScopedDevice('cpu'):
dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113)
cond_fn = lambda x: x.shape[0] not in {113, 1113}
dw_meta = _get_fields_with_cond(dw, cond_fn)
data_non_vmap = _get_non_vmap_data_fields()
return {k for k in dw_meta if k not in data_non_vmap}
return meta_fields - _get_non_vmap_data_fields()
with wp.ScopedDevice('cpu'):
mw = mjwarp.put_model(m)
cond_fn = lambda x: not hasattr(x, '_is_batched')
mw_meta = _get_fields_with_cond(mw, cond_fn)
if cls_name == 'Model':
return mw_meta
if cls_name == 'Option':
return {k[len('opt__') :] for k in mw_meta if k.startswith('opt')}
if cls_name == 'Statistic':
return {k[len('stat__') :] for k in mw_meta if k.startswith('stat')}
raise NotImplementedError(f'Unhandled class name {cls_name}.')
return meta_fields
def write_core_cls(
@@ -400,12 +404,7 @@ def write_core_cls(
extra_annotations: dict[str, type] | None = None,
):
"""Writes a core API class (e.g. Model/Data/Option/Statistic)."""
cls = {
'Model': mjwarp.Model,
'Data': mjwarp.Data,
'Option': mjwarp.Option,
'Statistic': mjwarp.Statistic,
}[cls_name]
cls = _CLS_MAP[cls_name]
annotations = dict(cls.__annotations__) # pytype: disable=attribute-error
if flatten_fields:
@@ -415,7 +414,7 @@ def write_core_cls(
for k, v in annotations.items():
if k not in meta_fields:
continue
if isinstance(v, wp.array) or type(v).__name__ == '_ArrayAnnotation':
if _is_array_spec(v):
annotations[k] = np.ndarray
warp_keys = annotations.keys()
@@ -454,43 +453,14 @@ def write_core_cls(
)
def _get_fields_with_cond(
d: Any,
cond_fn: Callable[[Any], bool],
s: Optional[Set[str]] = None,
prefix: str = '',
not_in: bool = False,
add_static: bool = False,
) -> Set[str]:
"""Recursively finds fields given a condition on the leading dimensions."""
if s is None:
s = set()
for f in dataclasses.fields(d):
attr = getattr(d, f.name)
if dataclasses.is_dataclass(f.type):
s = _get_fields_with_cond(attr, cond_fn, s, f.name + '__', not_in)
continue
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'):
continue
if cond_fn(attr):
s.add(prefix + f.name)
return s
def _get_non_vmap_data_fields() -> Set[str]:
"""Returns the fields that are not be vmapped but are still jax.Array."""
m = mujoco.MjModel.from_xml_string(_DUMMY_XML)
d = mujoco.MjData(m)
mujoco.mj_step(m, d)
dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113)
cond_fn = lambda x: x.shape[0] not in {113}
non_vmap = _get_fields_with_cond(dw, cond_fn, add_static=True)
return non_vmap
"""Returns Data fields that should not be vmapped."""
annotations = _get_annotations_recursive(dict(mjwarp.Data.__annotations__))
return {
name
for name, type_ in annotations.items()
if type_ in (int, float, bool) or _is_batched_field(type_) is False
}
def write_register_vmappable(target_fpath: epath.Path):
@@ -513,8 +483,7 @@ 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 _is_array_spec(wp_type):
return True
if wp_type in wp._src.types.value_types:
return True
@@ -535,14 +504,14 @@ def _to_jax_ndim(name: str, wp_type: Any) -> int:
def write_ndim_annotations(target_fpath: epath.Path):
"""Writes a dictionary of ndim for every warp field."""
ndim_annotations = {}
for cls in [mjwarp.Model, mjwarp.Data, mjwarp.Option, mjwarp.Statistic]:
ndim_annotations[cls.__name__] = {}
for cls_name, cls in _CLS_MAP.items():
ndim_annotations[cls_name] = {}
annotations = _get_annotations_recursive(cls.__annotations__)
for name, type_ in annotations.items():
if not _is_ffi_compatible(type_):
continue
ndim = _to_jax_ndim(name, type_)
ndim_annotations[cls.__name__][name] = ndim
ndim_annotations[cls_name][name] = ndim
with target_fpath.open('a') as f:
f.write('\n_NDIM = ' + _to_py_string(ndim_annotations))
@@ -550,36 +519,15 @@ def write_ndim_annotations(target_fpath: epath.Path):
def write_nworld_leading_dim(target_fpath: epath.Path):
"""Writes a dictionary of which MJW fields are batched."""
# TODO(btaba): check that batch fields have MJX jax.Array annotations, and
# that non-batch fields have np.ndarray annotations. Fail early.
m = mujoco.MjModel.from_xml_string(_DUMMY_XML)
d = mujoco.MjData(m)
batched = {}
for cls in [mjwarp.Model, mjwarp.Data, mjwarp.Option, mjwarp.Statistic]:
batched[cls.__name__] = {}
if cls.__name__ == 'Data':
with wp.ScopedDevice('cpu'):
dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113)
cond_fn = lambda x: x.shape[0] == 113
batched_fields = _get_fields_with_cond(dw, cond_fn)
else:
with wp.ScopedDevice('cpu'):
obj = mjwarp.put_model(m)
if cls.__name__ == 'Option':
obj = obj.opt
elif cls.__name__ == 'Statistic':
obj = obj.stat
cond_fn = lambda x: hasattr(x, '_is_batched')
batched_fields = _get_fields_with_cond(obj, cond_fn)
for cls_name, cls in _CLS_MAP.items():
batched[cls_name] = {}
all_annotations = _get_annotations_recursive(cls.__annotations__)
for name, type_ in all_annotations.items():
if not _is_ffi_compatible(type_):
continue
batched[cls.__name__][name] = name in batched_fields
# Same batch check as mujoco_warp._src.io._mark_batched.
batched[cls_name][name] = _is_batched_field(type_) is True
with target_fpath.open('a') as f:
f.write('\n_BATCH_DIM = ' + _to_py_string(batched))
mjx/mujoco/mjx/codegen/trace.py
+87 -73
View File
@@ -15,8 +15,8 @@
"""Static AST tracing to find MuJoCo Model and Data field usages."""
import ast
import builtins
import dataclasses
import functools
import importlib
import importlib.util
import os
@@ -26,6 +26,8 @@ from absl import logging
from etils import epath
from mujoco.mjx.codegen import file
_BUILTIN_NAMES = set(dir(builtins))
def _get_imported_module_names(fpath: epath.Path) -> Sequence[Tuple[str, str]]:
"""Returns set of (fully qualified module_name, alias) tuples."""
@@ -81,7 +83,7 @@ class FieldInfo:
class _FunctionFieldUsageVisitor(ast.NodeVisitor):
"""AST visitor to find attribute usages on 'm' and 'd' variables."""
"""AST visitor to find Model and Data field usages."""
def __init__(
self,
@@ -103,24 +105,28 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
self.generic_visit(node)
def add_field_usage(self, node: ast.Attribute, is_output: bool):
"""Adds field to the appropriate set."""
"""Adds field to the appropriate set based on mjwarp_field_info."""
attr_parts = []
curr_node = node
while isinstance(curr_node, ast.Attribute):
attr_parts.append(curr_node.attr)
curr_node = curr_node.value
if isinstance(curr_node, ast.Name):
attr_parts.reverse()
full_attribute_str = '__'.join(attr_parts)
in_field_info = full_attribute_str in self._mjwarp_field_info
if in_field_info:
if curr_node.id == 'm':
self.model_fields.add(full_attribute_str)
if curr_node.id == 'd':
self.data_fields.add(full_attribute_str)
if curr_node.id == 'd' and is_output:
self.data_out_fields.add(full_attribute_str)
if not isinstance(curr_node, ast.Name):
return
attr_parts.reverse()
full_attribute_str = '__'.join(attr_parts)
field_info = self._mjwarp_field_info.get(full_attribute_str)
if field_info is None:
return
# Classify by field_info.param_source; field names are guaranteed unique
# across Model and Data hierarchies by get_mjwarp_field_info.
if field_info.param_source == 'Model':
self.model_fields.add(full_attribute_str)
elif field_info.param_source == 'Data':
self.data_fields.add(full_attribute_str)
if is_output:
self.data_out_fields.add(full_attribute_str)
def visit_Attribute(self, node: ast.Attribute):
self.add_field_usage(node, self._in_outputs_context)
@@ -160,13 +166,14 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
"""Visit a function call node and recursively find all attribute usages."""
if isinstance(node.func, ast.Name):
called_fn_name = node.func.id
key = (hash(self._current_fpath), called_fn_name)
if key not in self._visited_fns:
self._visited_fns.add(key) # pyrefly: ignore[bad-argument-type]
next_fpath = self._module_fpaths.get(
called_fn_name, self._current_fpath
)
self.recurse_trace(next_fpath, called_fn_name)
if called_fn_name not in _BUILTIN_NAMES:
key = (self._current_fpath, called_fn_name)
if key not in self._visited_fns:
self._visited_fns.add(key)
next_fpath = self._module_fpaths.get(
called_fn_name, self._current_fpath
)
self.recurse_trace(next_fpath, called_fn_name)
elif isinstance(node.func, ast.Attribute):
parts = []
current = node.func
@@ -198,10 +205,10 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
self.add_field_usage(arg, is_output=True)
called_fn_name = '.'.join(parts[1:])
key = (hash(self._current_fpath), called_fn_name)
key = (self._current_fpath, called_fn_name)
next_fpath = self._module_fpaths.get(parts[0])
if next_fpath and key not in self._visited_fns:
self._visited_fns.add(key) # pyrefly: ignore[bad-argument-type]
self._visited_fns.add(key)
self.recurse_trace(next_fpath, called_fn_name)
self.generic_visit(node)
@@ -238,27 +245,6 @@ def trace_function(
if not target_fn_node:
raise ValueError(f'Function "{fn}" not found in "{fpath}".')
args_node = target_fn_node.args
args_tuple = tuple(
map(functools.partial(ast.get_source_segment, src), args_node.args)
)
check_args_tuple = (
('m: Model', 'd: Data'),
('m: types.Model', 'd: types.Data'),
('m', 'd'),
('m: Model',),
('m: types.Model',),
('m',),
)
if (
args_tuple[:2] not in check_args_tuple
and args_tuple[:1] not in check_args_tuple
):
raise ValueError(
f'Function "{fn}" in "{fpath}" must have arguments in'
f' {check_args_tuple} got {args_tuple}.'
)
if visited_fns is None:
visited_fns = set()
@@ -266,12 +252,22 @@ def trace_function(
for body in target_fn_node.body:
visitor.visit(body)
# Warn if tracing found no Model/Data field accesses.
if not visitor.model_fields and not visitor.data_fields:
logging.warning(
'Function "%s" in "%s" does not access any Model or Data fields.',
fn,
fpath,
)
# Detect RenderContext parameter independent of argument position.
render_context_in_caller = False
if len(target_fn_node.args.args) > 2:
third_param = target_fn_node.args.args[2]
if third_param.annotation:
annotation_str = ast.unparse(third_param.annotation)
render_context_in_caller = 'RenderContext' in annotation_str
for param in target_fn_node.args.args:
if param.annotation:
annotation_str = ast.unparse(param.annotation)
if 'RenderContext' in annotation_str:
render_context_in_caller = True
break
logging.info(
'End trace function "%s". Output fields: %s, RenderContext: %s',
@@ -285,7 +281,9 @@ def trace_function(
)
def get_mjwarp_field_info(src: str, get_cls_type_annotations) -> Dict[str, FieldInfo]:
def get_mjwarp_field_info(
src: str, get_cls_type_annotations
) -> Dict[str, FieldInfo]:
"""Return field info for mujoco_warp/_src/types.py."""
dataclass_map = {
'opt': 'Option',
@@ -293,33 +291,49 @@ def get_mjwarp_field_info(src: str, get_cls_type_annotations) -> Dict[str, Field
'efc': 'Constraint',
'contact': 'Contact',
}
field_info = {}
type_classes = get_cls_type_annotations(src)
for field, typ in type_classes['Model'].items():
if field == 'callback':
continue
if field in field_info:
raise AssertionError(f'Field {field} is duplicated in Model.')
if field in dataclass_map:
for sfield, styp in type_classes[dataclass_map[field]].items():
field_name = field + '__' + sfield
field_info[field_name] = FieldInfo('Model', styp, (1, field_name))
else:
field_info[field] = FieldInfo('Model', typ, (0, field))
for field, typ in type_classes['Data'].items():
if field in field_info:
raise AssertionError(f'Field {field} is duplicated.')
if field in dataclass_map:
for sfield, styp in type_classes[dataclass_map[field]].items():
field_name = field + '__' + sfield
field_info[field_name] = FieldInfo('Data', styp, (3, field_name))
else:
field_info[field] = FieldInfo('Data', typ, (2, field))
return field_info
def get_fields(cls_name, base_order, sub_order):
fields = {}
for field, typ in type_classes[cls_name].items():
if field == 'callback':
continue
if field in dataclass_map:
for sfield, styp in type_classes[dataclass_map[field]].items():
name = f'{field}__{sfield}'
fields[name] = FieldInfo(cls_name, styp, (sub_order, name))
else:
fields[field] = FieldInfo(cls_name, typ, (base_order, field))
return fields
model_fields = get_fields('Model', 0, 1)
data_fields = get_fields('Data', 2, 3)
# Field names must be uniquely named across Model and Data hierarchies because
# _FunctionFieldUsageVisitor.add_field_usage attributes AST attribute accesses
# to Model or Data based solely on the field name lookup in mjwarp_field_info
# (via field_info.param_source) rather than inspecting base variable names.
#
# Nested sub-dataclasses (Option, Statistic in Model; Constraint, Contact in
# Data) are flattened with container prefixes (e.g. 'opt__timestep',
# 'contact__pos'). This prefixing prevents collisions between sub-dataclasses
# that share generic field names (such as 'pos' or 'type' in Contact and
# Constraint). The collision check below guarantees that no direct or
# flattened sub-dataclass field name is shared between Model and Data.
#
# These classified fields (model_fields, data_fields, data_out_fields) are
# then used in generate_warp_shim to build the FFI shim argument signatures.
if collisions := model_fields.keys() & data_fields.keys():
raise ValueError(
f'Field name collisions in MuJoCo Warp types: {collisions}'
)
return {**model_fields, **data_fields}
def get_mjx_warp_field_info(src: str, get_cls_type_annotations) -> Dict[str, FieldInfo]:
def get_mjx_warp_field_info(
src: str, get_cls_type_annotations
) -> Dict[str, FieldInfo]:
"""Return field info for mjx/warp/types.py."""
field_info = {}
type_classes = get_cls_type_annotations(src)
mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
Vendored
+1
View File
@@ -41,6 +41,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import euler as euler
from mujoco.mjx.third_party.mujoco_warp._src.forward import forward as forward
from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_acceleration as fwd_acceleration
from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_actuation as fwd_actuation
from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_kinematics as fwd_kinematics
from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_position as fwd_position
from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_velocity as fwd_velocity
from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit
mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
Vendored
+152 -3
View File
@@ -18,6 +18,29 @@ from functools import lru_cache
import warp as wp
@wp.func
def solve_search_sums(grad: float, solution: float):
return wp.vec2(solution * solution, grad * solution)
@lru_cache(maxsize=None)
def _create_newton_decrement_func(matrix_size_static: int, vector_size_static: int):
@wp.func
def newton_decrement_func(
# In:
solution_tile: wp.tile[float, matrix_size_static, 1],
b: wp.array2d[float],
# Out:
search_out: wp.array2d[float],
):
grad_tile = wp.tile_load(b, shape=(vector_size_static, 1), offset=(0, 0), bounds_check=False)
active_solution = wp.tile_view(solution_tile, shape=(vector_size_static, 1), offset=(0, 0))
wp.tile_store(search_out, wp.tile_map(wp.mul, active_solution, -1.0), bounds_check=False)
return wp.tile_reduce(wp.add, wp.tile_map(solve_search_sums, grad_tile, active_solution))[0]
return newton_decrement_func
@lru_cache(maxsize=None)
def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
@wp.func
@@ -93,7 +116,109 @@ def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_st
@lru_cache(maxsize=None)
def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int):
def _create_blocked_cholesky_augmented_factorize_solve_func(
block_size: int,
matrix_size_static: int,
with_newton_decrement: bool,
vector_size_static: int,
):
WITH_NEWTON_DECREMENT = with_newton_decrement
border_size = block_size - 1
@wp.func
def blocked_cholesky_augmented_factorize_solve_func(
# In:
A: wp.array2d[float],
b: wp.array2d[float],
matrix_size: int,
# Out:
U_out: wp.array2d[float],
result_out: wp.array2d[float],
):
"""Factor A with b as an augmented border and reuse it as the forward solution."""
rhs_tile = wp.tile_zeros(shape=(matrix_size_static, 1), dtype=float, storage="shared")
for k in range(0, matrix_size, block_size):
end = k + block_size
input_rhs = wp.tile_load(b, shape=(block_size, 1), offset=(k, 0), storage="shared", bounds_check=False)
A_kk_tile = wp.tile_load(
A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
)
input_diagonal_rhs = wp.tile_view(input_rhs, shape=(border_size, 1), offset=(0, 0))
if end == matrix_size:
wp.tile_assign(A_kk_tile, input_diagonal_rhs, offset=(0, border_size))
# Keep the augmented rank-one correction below float32 precision.
A_kk_tile[border_size, border_size] = 1.0e30
for j in range(0, k, block_size):
U_block = wp.tile_load(
U_out, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
)
wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
diagonal_border = wp.tile_view(A_kk_tile, shape=(border_size, 1), offset=(0, border_size))
if end == matrix_size:
wp.tile_assign(rhs_tile, diagonal_border, offset=(k, 0))
wp.tile_store(U_out, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
for i in range(end, matrix_size, block_size):
A_ki_tile = wp.tile_load(
A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
)
if i + block_size == matrix_size:
wp.tile_assign(A_ki_tile, input_rhs, offset=(0, border_size))
for j in range(0, k, block_size):
U_jk_tile = wp.tile_load(
U_out, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
)
U_ji_tile = wp.tile_load(
U_out, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
panel_border = wp.tile_view(A_ki_tile, shape=(block_size, 1), offset=(0, border_size))
if i + block_size == matrix_size:
wp.tile_assign(rhs_tile, panel_border, offset=(k, 0))
wp.tile_store(U_out, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
for i in range(matrix_size - block_size, -1, -block_size):
i_end = i + block_size
tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(i_end, matrix_size, block_size):
U_tile = wp.tile_load(
U_out, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
)
x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
U_tile = wp.tile_load(
U_out, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_upper_solve_inplace(U_tile, tmp_tile)
sums = wp.vec2(0.0)
if wp.static(WITH_NEWTON_DECREMENT):
sums = wp.static(_create_newton_decrement_func(matrix_size_static, vector_size_static))(rhs_tile, b, result_out)
else:
wp.tile_store(result_out, rhs_tile, offset=(0, 0), bounds_check=False)
return sums
return blocked_cholesky_augmented_factorize_solve_func
@lru_cache(maxsize=None)
def _create_blocked_cholesky_solve_func(
block_size: int,
matrix_size_static: int,
with_newton_decrement: bool,
vector_size_static: int,
):
WITH_NEWTON_DECREMENT = with_newton_decrement
@wp.func
def blocked_cholesky_solve_func(
# In:
@@ -101,7 +226,7 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
b: wp.array2d[float],
matrix_size: int,
# Out:
x: wp.array2d[float],
result_out: wp.array2d[float],
):
"""Block Cholesky solve.
@@ -141,6 +266,30 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
wp.tile_upper_solve_inplace(U_tile, tmp_tile)
wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
sums = wp.vec2(0.0)
if wp.static(WITH_NEWTON_DECREMENT):
sums = wp.static(_create_newton_decrement_func(matrix_size_static, vector_size_static))(rhs_tile, b, result_out)
else:
wp.tile_store(result_out, rhs_tile, offset=(0, 0), bounds_check=False)
return sums
return blocked_cholesky_solve_func
def create_blocked_cholesky_augmented_factorize_solve_func(block_size: int, matrix_size_static: int):
return _create_blocked_cholesky_augmented_factorize_solve_func(block_size, matrix_size_static, False, 0)
def create_blocked_cholesky_augmented_factorize_solve_newton_func(
block_size: int, matrix_size_static: int, vector_size_static: int
):
return _create_blocked_cholesky_augmented_factorize_solve_func(block_size, matrix_size_static, True, vector_size_static)
def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int):
return _create_blocked_cholesky_solve_func(block_size, matrix_size_static, False, 0)
def create_blocked_cholesky_solve_newton_func(block_size: int, matrix_size_static: int, vector_size_static: int):
return _create_blocked_cholesky_solve_func(block_size, matrix_size_static, True, vector_size_static)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
Vendored
+1 -1
View File
@@ -28,7 +28,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@event_scope
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+9 -3
View File
@@ -52,7 +52,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
# TODO(team): improve compile time to enable backward pass
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float)
mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
@@ -172,7 +172,7 @@ def ccd_hfield_kernel_builder(
"""Kernel builder for heightfield CCD collisions (no multiccd args)."""
# runs convex collision on a set of geom pairs to recover contact info
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def ccd_hfield_kernel(
# Model:
opt_ccd_tolerance: wp.array[float],
@@ -474,6 +474,9 @@ def ccd_hfield_kernel_builder(
epa_pr,
epa_norm2,
epa_horizon,
wp.static(warn_overflow),
worldid,
overflow_out,
)
if ncontact == 0:
@@ -833,6 +836,9 @@ def ccd_kernel_builder(
epa_pr_in[ccdid],
epa_norm2_in[ccdid],
epa_horizon_in[ccdid],
wp.static(warn_overflow),
worldid,
overflow_out,
)
if dist >= gap and not is_collision_sensor:
@@ -946,7 +952,7 @@ def ccd_kernel_builder(
)
# runs convex collision on a set of geom pairs to recover contact info (non-heightfield)
@wp.kernel(module="unique", enable_backward=False, launch_bounds=(block_dim, _CCD_MIN_BLOCKS))
@wp.kernel(module="unique", enable_backward=False, grid_stride=False, launch_bounds=(block_dim, _CCD_MIN_BLOCKS))
def ccd_kernel(
# Model:
opt_ccd_tolerance: wp.array[float],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+1 -1
View File
@@ -28,7 +28,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.struct
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+5 -5
View File
@@ -41,7 +41,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
# Corresponding table to MuJoCo's mjCOLLISIONFUNC table in engine_collision_driver.c
MJ_COLLISION_TABLE = {
@@ -373,7 +373,7 @@ def _add_geom_pair(
@cache_kernel
def _sap_project(opt_broadphase: int):
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def sap_project(
# Model:
ngeom: int,
@@ -430,7 +430,7 @@ def _sap_broadphase(
enable_sleep: bool = False,
incremental: bool = False,
):
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
ngeom: int,
@@ -538,7 +538,7 @@ def _sap_broadphase(
@cache_kernel
def _segmented_sort(tile_size: int):
@wp.kernel(module="unique")
@wp.kernel(module="unique", grid_stride=False)
def segmented_sort(
# In:
projection_lower_in: wp.array2d[float],
@@ -691,7 +691,7 @@ def _nxn_broadphase(
enable_sleep: bool = False,
incremental: bool = False,
):
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
geom_type: wp.array[int],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py
Vendored
+22 -4
View File
@@ -36,7 +36,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import OverflowType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.func
@@ -796,6 +796,7 @@ def _collide_mesh_triangle(
epa_horizon: wp.array[int],
tolerance: float,
ccd_iterations: int,
warn_overflow: bool,
# Data out:
overflow_out: wp.array[int],
# Out:
@@ -861,6 +862,9 @@ def _collide_mesh_triangle(
epa_pr,
epa_norm2,
epa_horizon,
warn_overflow,
worldid,
overflow_out,
)
if ncontact > 0 and dist < margin + tri_radius:
@@ -1285,7 +1289,7 @@ def _plane_vertex(
return True, dist, contact_pos, nrm_out
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def _flex_internal_collisions_detect(
# Model:
nflex: int,
@@ -1394,7 +1398,7 @@ def _flex_internal_collisions_detect(
)
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def _flex_tet_internal_collisions_detect(
# Model:
nflex: int,
@@ -1840,6 +1844,7 @@ def _flex_selfcollision_narrowphase(
# Model:
nflex: int,
opt_ccd_tolerance: wp.array[float],
opt_warn_overflow: bool,
flex_dim: wp.array[int],
flex_vertadr: wp.array[int],
flex_elemadr: wp.array[int],
@@ -2015,6 +2020,9 @@ def _flex_selfcollision_narrowphase(
epa_pr_out[pairid],
epa_norm2_out[pairid],
epa_horizon_out[pairid],
opt_warn_overflow,
worldid,
overflow_out,
)
phys_dist = dist
@@ -2055,11 +2063,12 @@ def _flex_selfcollision_narrowphase(
)
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def _flex_active_element_collisions_detect(
# Model:
nflex: int,
opt_ccd_tolerance: wp.array[float],
opt_warn_overflow: bool,
flex_selfcollide: wp.array[int],
flex_dim: wp.array[int],
flex_vertadr: wp.array[int],
@@ -2234,6 +2243,9 @@ def _flex_active_element_collisions_detect(
epa_pr_out[unique_thread_id],
epa_norm2_out[unique_thread_id],
epa_horizon_out[unique_thread_id],
opt_warn_overflow,
worldid,
overflow_out,
)
phys_dist = dist
@@ -2445,6 +2457,7 @@ def _flex_narrowphase_unified(
epa_horizon[ccdid],
tolerance,
ccd_iterations,
opt_warn_overflow,
overflow_out,
cand_dist_out,
cand_pos_out,
@@ -2507,6 +2520,9 @@ def _flex_narrowphase_unified(
epa_pr[ccdid],
epa_norm2[ccdid],
epa_horizon[ccdid],
opt_warn_overflow,
worldid,
overflow_out,
)
if ncontact > 0 and dist < margin + tri_radius:
@@ -3650,6 +3666,7 @@ def flex_collision(m: Model, d: Data, ctx):
inputs=[
m.nflex,
m.opt.ccd_tolerance,
m.opt.warn_overflow,
m.flex_dim,
m.flex_vertadr,
m.flex_elemadr,
@@ -3716,6 +3733,7 @@ def flex_collision(m: Model, d: Data, ctx):
inputs=[
m.nflex,
m.opt.ccd_tolerance,
m.opt.warn_overflow,
m.flex_selfcollide,
m.flex_dim,
m.flex_vertadr,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+133 -44
View File
@@ -20,6 +20,7 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import OverflowType
from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
@@ -67,6 +68,8 @@ class GJKResult:
simplex2: mat43
simplex_index1: wp.vec4i
simplex_index2: wp.vec4i
index1: int
index2: int
@wp.struct
@@ -170,6 +173,7 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
edge_localid = geom.graphadr + 2 + 2 * numvert
prev = int(-1)
imax = wp.where(geom.index > -1, geom.index, 0)
max_dist = wp.dot(local_dir, geom.vert[geom.vertadr + geom.graph[vert_globalid + imax]])
# hillclimb until no change
while imax != prev:
@@ -272,11 +276,6 @@ def _linear_combine(n: int, scl: wp.vec4, mat: mat43) -> wp.vec3:
return scl[0] * mat[0] + scl[1] * mat[1] + scl[2] * mat[2] + scl[3] * mat[3]
@wp.func
def _almost_equal(v1: wp.vec3, v2: wp.vec3) -> bool:
return wp.abs(v1[0] - v2[0]) < MINVAL and wp.abs(v1[1] - v2[1]) < MINVAL and wp.abs(v1[2] - v2[2]) < MINVAL
@wp.func
def _subdistance(n: int, simplex: mat43) -> wp.vec4:
if n == 4:
@@ -592,6 +591,45 @@ def _S1D(s1: wp.vec3, s2: wp.vec3) -> wp.vec2:
return wp.vec2(0.0, 1.0)
@wp.func
def _gjk_support(
# In:
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
x_k: wp.vec3,
x_norm: float,
simplex: mat43,
n: int,
is_discrete: bool,
) -> Tuple[SupportPoint, SupportPoint]:
dir_neg = x_k / x_norm
# tuning for discrete geoms when direction is noisy
if is_discrete and x_norm < 1e-4:
if n == 2:
edge = simplex[1] - simplex[0]
edge_norm2 = wp.dot(edge, edge)
if edge_norm2 > MINVAL2:
proj = wp.dot(dir_neg, edge) / edge_norm2
dir_neg = dir_neg - proj * edge
dir_norm = wp.length(dir_neg)
if dir_norm > MINVAL:
dir_neg = dir_neg / dir_norm
elif n == 3:
e1 = simplex[1] - simplex[0]
e2 = simplex[2] - simplex[0]
normal = wp.cross(e1, e2)
normal_norm = wp.length(normal)
if normal_norm > MINVAL:
dir_neg = wp.sign(wp.dot(dir_neg, normal)) * normal / normal_norm
sp1 = support(geom1, geomtype1, -dir_neg)
sp2 = support(geom2, geomtype2, dir_neg)
return sp1, sp2
@wp.func
def gjk(
# In:
@@ -614,33 +652,34 @@ def gjk(
simplex_index1 = wp.vec4i()
simplex_index2 = wp.vec4i()
n = int(0)
lmbda = wp.vec4() # barycentric coordinates
tol2 = tolerance * tolerance
epsilon = wp.where(is_discrete, 0.0, 0.5 * tol2)
lmbda = wp.vec4(1.0, 0.0, 0.0, 0.0) # barycentric coordinates
# for discrete geoms GJK is guaranteed to converge in a finite number of iterations
# so we can ignore tolerance
# TODO(kbayes): look into relative tolerances based off of xnorm
epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance)
min_norm = wp.where(is_discrete, MINVAL, tolerance)
min_tol = wp.where(is_discrete, MINVAL, tolerance)
# set initial guess
x_k = x1_0 - x2_0
xnorm2_old = FLOAT_MAX
xnorm2 = wp.dot(x_k, x_k)
xnorm = wp.sqrt(xnorm2)
xnorm_prev = float(0.0)
for _ in range(gjk_iterations):
xnorm2 = wp.dot(x_k, x_k)
# TODO(kbayes): determine new constant here
if xnorm2 < tol2 or wp.abs(xnorm2_old - xnorm2) < tol2:
if xnorm < min_norm or wp.abs(xnorm_prev - xnorm) < min_tol:
break
xnorm2_old = xnorm2
dir_neg = x_k / wp.sqrt(xnorm2)
# compute kth support point in geom1
sp = support(geom1, geomtype1, -dir_neg)
simplex1[n] = sp.point
geom1.index = sp.cached_index
simplex_index1[n] = sp.vertex_index
# compute the support point with direction tuning
sp1, sp2 = _gjk_support(geom1, geom2, geomtype1, geomtype2, x_k, xnorm, simplex, n, is_discrete)
simplex1[n] = sp1.point
geom1.index = sp1.cached_index
simplex_index1[n] = sp1.vertex_index
# compute kth support point in geom2
sp = support(geom2, geomtype2, dir_neg)
simplex2[n] = sp.point
geom2.index = sp.cached_index
simplex_index2[n] = sp.vertex_index
simplex2[n] = sp2.point
geom2.index = sp2.cached_index
simplex_index2[n] = sp2.vertex_index
# compute the kth support point
simplex[n] = simplex1[n] - simplex2[n]
@@ -655,6 +694,8 @@ def gjk(
result = GJKResult()
result.dim = 0
result.dist = FLOAT_MAX
result.index1 = geom1.index
result.index2 = geom2.index
return result
elif cutoff < FLOAT_MAX:
vs = wp.dot(x_k, simplex[n])
@@ -662,6 +703,8 @@ def gjk(
result = GJKResult()
result.dim = 0
result.dist = FLOAT_MAX
result.index1 = geom1.index
result.index2 = geom2.index
return result
# run the distance subalgorithm to compute the barycentric coordinates
@@ -686,20 +729,17 @@ def gjk(
if n < 1:
break
# get the next iteration of x_k
x_next = _linear_combine(n, lmbda, simplex)
# x_k has converged to minimum
if _almost_equal(x_next, x_k):
break
# copy next iteration into x_k
x_k = x_next
# we have a tetrahedron containing the origin so return early
if n == 4:
xnorm = 0.0
break
# get the next iteration of x_k
x_k = _linear_combine(n, lmbda, simplex)
xnorm_prev = xnorm
xnorm2 = wp.dot(x_k, x_k)
xnorm = wp.sqrt(xnorm2)
result = GJKResult()
# compute the approximate witness points
@@ -707,7 +747,7 @@ def gjk(
# are the witness points
result.x1 = wp.where(n == 0, x1_0, _linear_combine(n, lmbda, simplex1))
result.x2 = wp.where(n == 0, x2_0, _linear_combine(n, lmbda, simplex2))
result.dist = wp.norm_l2(x_k)
result.dist = xnorm
result.dim = n
result.simplex1 = simplex1
@@ -715,6 +755,8 @@ def gjk(
result.simplex_index1 = simplex_index1
result.simplex_index2 = simplex_index2
result.simplex = simplex
result.index1 = geom1.index
result.index2 = geom2.index
return result
@@ -1182,26 +1224,48 @@ def _polytope4(
pt.vert_index[6] = simplex_index1[3]
pt.vert_index[7] = simplex_index2[3]
dist = wp.vec4()
idx = int(0)
# if the origin is on a face, replace the 3-simplex with a 2-simplex
if _attach_face(pt, 0, 0, 1, 2) < MIN_DIST4:
dist[0] = _attach_face(pt, 0, 0, 1, 2)
if dist[0] < MIN_DIST4:
pt.status = -1
return pt, _replace_simplex3(pt, 0, 1, 2)
if _attach_face(pt, 1, 0, 3, 1) < MIN_DIST4:
dist[1] = _attach_face(pt, 1, 0, 3, 1)
if dist[1] < MIN_DIST4:
pt.status = -1
return pt, _replace_simplex3(pt, 0, 3, 1)
idx = wp.where(dist[0] < dist[1], 0, 1)
if _attach_face(pt, 2, 0, 2, 3) < MIN_DIST4:
dist[2] = _attach_face(pt, 2, 0, 2, 3)
if dist[2] < MIN_DIST4:
pt.status = -1
return pt, _replace_simplex3(pt, 0, 2, 3)
idx = wp.where(dist[2] < dist[idx], 2, idx)
if _attach_face(pt, 3, 3, 2, 1) < MIN_DIST4:
dist[3] = _attach_face(pt, 3, 3, 2, 1)
if dist[3] < MIN_DIST4:
pt.status = -1
return pt, _replace_simplex3(pt, 3, 2, 1)
idx = wp.where(dist[3] < dist[idx], 3, idx)
if not _test_tetra(simplex[0], simplex[1], simplex[2], simplex[3]):
pt.status = 12
return pt, GJKResult()
if dist[idx] > MINVAL:
pt.status = 12
return pt, GJKResult()
# fallback to closest face
pt.status = -1
if idx == 0:
return pt, _replace_simplex3(pt, 0, 1, 2)
elif idx == 1:
return pt, _replace_simplex3(pt, 0, 3, 1)
elif idx == 2:
return pt, _replace_simplex3(pt, 0, 2, 3)
else:
return pt, _replace_simplex3(pt, 3, 2, 1)
# set polytope counts
pt.nvert = 4
@@ -1251,6 +1315,10 @@ def _epa(
geomtype1: int,
geomtype2: int,
is_discrete: bool,
warn_overflow: bool,
worldid: int,
# Data out:
overflow_out: wp.array[int],
) -> Tuple[float, wp.vec3, wp.vec3, int]:
"""Recover penetration data from two geoms in contact given an initial polytope."""
upper = FLOAT_MAX
@@ -1320,7 +1388,9 @@ def _epa(
pt.nhorizon = _add_edge(pt, face[1], face[2])
pt.nhorizon = _add_edge(pt, face[2], face[0])
if pt.nhorizon == -1:
wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
if warn_overflow:
wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
wp.atomic_or(overflow_out, worldid, OverflowType.EPA_HORIZON)
idx = -1
break
@@ -1337,7 +1407,9 @@ def _epa(
pt.nhorizon = _add_edge(pt, face[1], face[2])
pt.nhorizon = _add_edge(pt, face[2], face[0])
if pt.nhorizon == -1:
wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
if warn_overflow:
wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
wp.atomic_or(overflow_out, worldid, OverflowType.EPA_HORIZON)
idx = -1
break
@@ -2302,6 +2374,8 @@ def gjk_phase(
if size1 + size2 > 0.0:
cutoff += full_margin1 + full_margin2
result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
geom1.index = result.index1
geom2.index = result.index2
# shallow penetration, inflate contact
if result.dist > tolerance:
@@ -2318,6 +2392,8 @@ def gjk_phase(
cutoff -= full_margin1 + full_margin2
result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
geom1.index = result.index1
geom2.index = result.index2
# no penetration depth to recover
if result.dist > tolerance or result.dim < 2:
@@ -2342,6 +2418,10 @@ def epa_phase(
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
warn_overflow: bool,
worldid: int,
# Data out:
overflow_out: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
pt = Polytope()
@@ -2413,7 +2493,9 @@ def epa_phase(
return result.dist, 1, result.x1, result.x2, -1
is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
dist, x1, x2, idx = _epa(
tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete, warn_overflow, worldid, overflow_out
)
if idx == -1:
return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
@@ -2447,6 +2529,10 @@ def ccd(
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
warn_overflow: bool,
worldid: int,
# Data out:
overflow_out: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""General convex collision detection via GJK/EPA."""
needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
@@ -2468,4 +2554,7 @@ def ccd(
face_pr,
face_norm2,
horizon,
warn_overflow,
worldid,
overflow_out,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py
Vendored
+7 -66
View File
@@ -13,79 +13,20 @@
# limitations under the License.
# ==============================================================================
from typing import Any, Tuple
from typing import Tuple
import warp as wp
MJ_MINVAL = 1e-15
MJ_MAXVAL = 1e10
from mujoco.mjx.third_party.mujoco_warp._src.math import closest_segment_point
from mujoco.mjx.third_party.mujoco_warp._src.math import closest_segment_to_segment_points
from mujoco.mjx.third_party.mujoco_warp._src.math import normalize_with_norm
from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
wp.set_module_options({"enable_backward": False})
@wp.func
def safe_div(x: Any, y: Any) -> Any:
return x / wp.where(y != 0.0, y, MJ_MINVAL)
@wp.func
def normalize_with_norm(x: Any):
norm = wp.length(x)
if norm == 0.0:
return x, 0.0
return x / norm, norm
@wp.func
def closest_segment_point(a: wp.vec3, b: wp.vec3, pt: wp.vec3) -> wp.vec3:
"""Returns the closest point on the a-b line segment to a point pt."""
ab = b - a
t = wp.dot(pt - a, ab) / (wp.dot(ab, ab) + 1e-6)
return a + wp.clamp(t, 0.0, 1.0) * ab
@wp.func
def closest_segment_point_and_dist(a: wp.vec3, b: wp.vec3, pt: wp.vec3) -> Tuple[wp.vec3, float]:
"""Returns closest point on the line segment and the distance squared."""
closest = closest_segment_point(a, b, pt)
dist = wp.dot((pt - closest), (pt - closest))
return closest, dist
@wp.func
def closest_segment_to_segment_points(a0: wp.vec3, a1: wp.vec3, b0: wp.vec3, b1: wp.vec3) -> Tuple[wp.vec3, wp.vec3]:
"""Returns closest points between two line segments."""
dir_a, len_a = normalize_with_norm(a1 - a0)
dir_b, len_b = normalize_with_norm(b1 - b0)
half_len_a = len_a * 0.5
half_len_b = len_b * 0.5
a_mid = a0 + dir_a * half_len_a
b_mid = b0 + dir_b * half_len_b
trans = a_mid - b_mid
dira_dot_dirb = wp.dot(dir_a, dir_b)
dira_dot_trans = wp.dot(dir_a, trans)
dirb_dot_trans = wp.dot(dir_b, trans)
denom = 1.0 - dira_dot_dirb * dira_dot_dirb
orig_t_a = (-dira_dot_trans + dira_dot_dirb * dirb_dot_trans) / (denom + 1e-6)
orig_t_b = dirb_dot_trans + orig_t_a * dira_dot_dirb
t_a = wp.clamp(orig_t_a, -half_len_a, half_len_a)
t_b = wp.clamp(orig_t_b, -half_len_b, half_len_b)
best_a = a_mid + dir_a * t_a
best_b = b_mid + dir_b * t_b
new_a, d1 = closest_segment_point_and_dist(a0, a1, best_b)
new_b, d2 = closest_segment_point_and_dist(b0, b1, best_a)
if d1 < d2:
return new_a, best_b
return best_a, new_b
class vec8f(wp.types.vector(length=8, dtype=wp.float32)):
pass
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+1 -1
View File
@@ -34,7 +34,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec_pluginattr
from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.struct
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+450 -455
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+11 -12
View File
@@ -29,10 +29,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.kernel
@@ -44,14 +43,14 @@ def _qderiv_actuator_passive_vel(
actuator_biastype: wp.array[int],
actuator_actadr: wp.array[int],
actuator_actnum: wp.array[int],
actuator_forcelimited: wp.array[bool],
actuator_dynprm: wp.array2d[vec10],
actuator_gainprm: wp.array2d[vec10],
actuator_biasprm: wp.array2d[vec10],
actuator_actlimited: wp.array[bool],
actuator_dynprm: wp.array2d[vec10f],
actuator_gainprm: wp.array2d[vec10f],
actuator_biasprm: wp.array2d[vec10f],
actuator_actearly: wp.array[bool],
actuator_forcerange: wp.array2d[wp.vec2],
actuator_actrange: wp.array2d[wp.vec2],
actuator_actearly: wp.array[bool],
actuator_forcelimited: wp.array[bool],
actuator_forcerange: wp.array2d[wp.vec2],
# Data in:
act_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
@@ -1141,14 +1140,14 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
m.actuator_biastype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_forcelimited,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_gainprm,
m.actuator_biasprm,
m.actuator_actearly,
m.actuator_forcerange,
m.actuator_actlimited,
m.actuator_actrange,
m.actuator_actearly,
m.actuator_forcelimited,
m.actuator_forcerange,
d.act,
d.ctrl,
d.act_dot,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+55 -44
View File
@@ -43,7 +43,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import OverflowType
from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -138,10 +138,10 @@ def _next_activation(
actuator_dyntype: wp.array[int],
actuator_actadr: wp.array[int],
actuator_actnum: wp.array[int],
actuator_dynprm: wp.array2d[vec10],
actuator_gainprm: wp.array2d[vec10],
actuator_biasprm: wp.array2d[vec10],
actuator_actlimited: wp.array[bool],
actuator_dynprm: wp.array2d[vec10f],
actuator_gainprm: wp.array2d[vec10f],
actuator_biasprm: wp.array2d[vec10f],
actuator_actrange: wp.array2d[wp.vec2],
# Data in:
act_in: wp.array2d[float],
@@ -286,10 +286,10 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
m.actuator_dyntype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_gainprm,
m.actuator_biasprm,
m.actuator_actlimited,
m.actuator_actrange,
d.act,
d.act_dot,
@@ -451,10 +451,10 @@ def _rk_perturb_state(
m.actuator_dyntype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_gainprm,
m.actuator_biasprm,
m.actuator_actlimited,
m.actuator_actrange,
act_t0,
d.act_dot,
@@ -613,13 +613,12 @@ def implicit(m: Model, d: Data):
@event_scope
def fwd_position(m: Model, d: Data, factorize: bool = True):
"""Position-dependent computations.
def fwd_kinematics(m: Model, d: Data):
"""Kinematics-dependent computations.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
factorize: Flag to factorize interia matrix.
"""
smooth.kinematics(m, d)
smooth.com_pos(m, d)
@@ -628,11 +627,24 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
smooth.tendon(m, d)
sleep_enabled = bool(m.opt.enableflags & EnableBit.SLEEP) and not bool(m.opt.disableflags & DisableBit.ISLAND)
if sleep_enabled and m.ntendon > 0:
sleep.wake_tendon(m, d)
sleep.update_sleep_trees(m, d)
@event_scope
def fwd_position(m: Model, d: Data, factorize: bool = True):
"""Position-dependent computations.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
factorize: Flag to factorize inertia matrix.
"""
fwd_kinematics(m, d)
sleep_enabled = bool(m.opt.enableflags & EnableBit.SLEEP) and not bool(m.opt.disableflags & DisableBit.ISLAND)
smooth.crb(m, d)
smooth.tendon_armature(m, d)
if factorize:
@@ -751,16 +763,16 @@ def _actuator_force(
actuator_biastype: wp.array[int],
actuator_actadr: wp.array[int],
actuator_actnum: wp.array[int],
actuator_ctrllimited: wp.array[bool],
actuator_forcelimited: wp.array[bool],
actuator_dynprm: wp.array2d[vec10],
actuator_gainprm: wp.array2d[vec10],
actuator_biasprm: wp.array2d[vec10],
actuator_actlimited: wp.array[bool],
actuator_dynprm: wp.array2d[vec10f],
actuator_gainprm: wp.array2d[vec10f],
actuator_biasprm: wp.array2d[vec10f],
actuator_actearly: wp.array[bool],
actuator_ctrlrange: wp.array2d[wp.vec2],
actuator_forcerange: wp.array2d[wp.vec2],
actuator_actrange: wp.array2d[wp.vec2],
actuator_actearly: wp.array[bool],
actuator_forcelimited: wp.array[bool],
actuator_forcerange: wp.array2d[wp.vec2],
actuator_ctrllimited: wp.array[bool],
actuator_ctrlrange: wp.array2d[wp.vec2],
actuator_acc0: wp.array2d[float],
actuator_lengthrange: wp.array2d[wp.vec2],
# Data in:
@@ -819,6 +831,7 @@ def _actuator_force(
adr += 1
# integral
x_I = 0.0
if slots[1] >= 0:
x_I = act_in[worldid, adr]
input_mode = int(gainprm[8])
@@ -1163,16 +1176,16 @@ def fwd_actuation(m: Model, d: Data):
m.actuator_biastype,
m.actuator_actadr,
m.actuator_actnum,
m.actuator_ctrllimited,
m.actuator_forcelimited,
m.actuator_actlimited,
m.actuator_dynprm,
m.actuator_gainprm,
m.actuator_biasprm,
m.actuator_actearly,
m.actuator_ctrlrange,
m.actuator_forcerange,
m.actuator_actlimited,
m.actuator_actrange,
m.actuator_actearly,
m.actuator_forcelimited,
m.actuator_forcerange,
m.actuator_ctrllimited,
m.actuator_ctrlrange,
m.actuator_acc0,
m.actuator_lengthrange,
d.act,
@@ -1311,6 +1324,20 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
def _energy_pos(m: Model, d: Data):
if m.opt.enableflags & EnableBit.ENERGY:
if m.sensor_e_potential == 0: # not computed by sensor
sensor.energy_pos(m, d)
else:
d.energy.zero_()
def _energy_vel(m: Model, d: Data):
if m.opt.enableflags & EnableBit.ENERGY:
if m.sensor_e_kinetic == 0: # not computed by sensor
sensor.energy_vel(m, d)
@event_scope
def forward(m: Model, d: Data):
"""Forward dynamics."""
@@ -1319,23 +1346,14 @@ def forward(m: Model, d: Data):
sleep.wake(m, d)
sleep.update_sleep(m, d)
energy = m.opt.enableflags & EnableBit.ENERGY
fwd_position(m, d, factorize=False)
d.sensordata.zero_()
sensor.sensor_pos(m, d)
if energy:
if m.sensor_e_potential == 0: # not computed by sensor
sensor.energy_pos(m, d)
else:
d.energy.zero_()
_energy_pos(m, d)
fwd_velocity(m, d)
sensor.sensor_vel(m, d)
if energy:
if m.sensor_e_kinetic == 0: # not computed by sensor
sensor.energy_vel(m, d)
_energy_vel(m, d)
if not (m.opt.disableflags & DisableBit.ACTUATION):
if m.callback.control:
@@ -1365,23 +1383,16 @@ def step(m: Model, d: Data):
@event_scope
def step1(m: Model, d: Data):
"""Advance simulation in two phases: before input is set by user."""
energy = m.opt.enableflags & EnableBit.ENERGY
fwd_position(m, d)
d.sensordata.zero_()
sensor.sensor_pos(m, d)
if energy:
if m.sensor_e_potential == 0: # not computed by sensor
sensor.energy_pos(m, d)
else:
d.energy.zero_()
_energy_pos(m, d)
fwd_velocity(m, d)
sensor.sensor_vel(m, d)
if energy:
if m.sensor_e_kinetic == 0: # not computed by sensor
sensor.energy_vel(m, d)
_energy_vel(m, d)
if not (m.opt.disableflags & DisableBit.ACTUATION):
if m.callback.control:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
Vendored
+1 -1
View File
@@ -20,7 +20,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.func
mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
Vendored
+16 -1
View File
@@ -29,7 +29,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.kernel
@@ -119,6 +119,13 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
smooth.solve_m(m, d, qacc, qfrc)
@wp.kernel
def _zero_qfrc_constraint_nefc(nefc_in: wp.array[int], qfrc_constraint_out: wp.array2d[float]):
worldid, dofid = wp.tid()
if nefc_in[worldid] == 0:
qfrc_constraint_out[worldid, dofid] = 0.0
def inv_constraint(m: Model, d: Data):
"""Inverse constraint solver."""
# no constraints
@@ -126,6 +133,14 @@ def inv_constraint(m: Model, d: Data):
d.qfrc_constraint.zero_()
return
if m.is_sparse:
wp.launch(
_zero_qfrc_constraint_nefc,
dim=(d.nworld, m.nv),
inputs=[d.nefc],
outputs=[d.qfrc_constraint],
)
ctx = solver.create_inverse_context(m, d)
solver.init_context(m, d, ctx, grad=False)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+379 -159
View File
@@ -26,13 +26,17 @@ from mujoco.mjx.third_party.mujoco_warp._src import math as mjmath
from mujoco.mjx.third_party.mujoco_warp._src import render_util
from mujoco.mjx.third_party.mujoco_warp._src import sleep
from mujoco.mjx.third_party.mujoco_warp._src import smooth
from mujoco.mjx.third_party.mujoco_warp._src import support
from mujoco.mjx.third_party.mujoco_warp._src import types
from mujoco.mjx.third_party.mujoco_warp._src import warp_util
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import MJ_COLLISION_TABLE
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
wp.set_module_options({"default_grid_stride": False})
def _is_array_spec(typ) -> bool:
"""Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
@@ -70,7 +74,12 @@ def _create_array(data: Any, spec, sizes: dict[str, int], batch_size: int = 1) -
return None
return wp.array(np.array(data), dtype=spec.dtype)
shape = tuple(batch_size if dim == "*" else (sizes[dim] if isinstance(dim, str) else dim) for dim in spec_shape)
shape = tuple(
batch_size
if dim == "*"
else (int(dim) if isinstance(dim, str) and dim.isdigit() else (sizes[dim] if isinstance(dim, str) else dim))
for dim in spec_shape
)
is_batched = spec_shape[0] in ("*", "nworld")
@@ -94,14 +103,13 @@ def _create_constraint(
mjm,
nworld: int,
njmax: int,
njmax_nnz: int,
sizes: dict,
mjd=None,
) -> types.Constraint:
"""Construct a types.Constraint with standard and island local fields allocated properly."""
efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
sparse = is_sparse(mjm)
# The JTDAJ block list is only consumed by the sparse Newton Hessian assembly (_JTDAJ_sparse).
# The JTDAJ block list is only consumed by the sparse Newton Hessian assembly (_JTDACJ_sparse).
jtdaj_active = sparse and mjm.opt.solver == mujoco.mjtSolver.mjSOL_NEWTON
for f in dataclasses.fields(types.Constraint):
@@ -145,18 +153,31 @@ def _jtdaj_groups(mjd: mujoco.MjData) -> tuple[np.ndarray, np.ndarray]:
def _get_nflexintcell(mjm: mujoco.MjModel) -> int:
nflexintcell = 0
if mjm.nflex > 0 and hasattr(mjm, "flex_interp"):
if mjm.nflex > 0:
for fi in range(mjm.nflex):
order = abs(int(mjm.flex_interp[fi]))
if order == 0:
continue
if hasattr(mjm, "flex_edgeequality") and mjm.flex_edgeequality[fi] == 3:
if mjm.flex_edgeequality[fi] == 3:
continue
cx, cy, cz = mjm.flex_cellnum[fi]
nflexintcell += int(cx) * int(cy) * int(cz)
return nflexintcell
def _get_nflexface(mjm: mujoco.MjModel) -> int:
nflexface = 0
if mjm.nflex > 0:
for fi in range(mjm.nflex):
order = mjm.flex_interp[fi]
if order >= 0:
continue
cx, cy, cz = mjm.flex_cellnum[fi]
nfaces = 2 * (cy * cz + cx * cz + cx * cy)
nflexface += int(nfaces)
return nflexface
def is_sparse(mjm: mujoco.MjModel) -> bool:
if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
if mjm.nv > 32:
@@ -177,62 +198,44 @@ def _m_blocks(mjm: mujoco.MjModel):
return [(int(adr), int(num)) for adr, num in zip(mjm.tree_dofadr, mjm.tree_dofnum) if num > 0]
def _m_allow_dense(mjm: mujoco.MjModel) -> bool:
"""Whether any block may use the packed dense layout (tendon armature forces all-sparse)."""
# tendon armature accumulates into M in CSR layout, which the packed block layout cannot represent
return not (mjm.ntendon and np.any(mjm.tendon_armature))
def m_block_layout(mjm: mujoco.MjModel) -> dict:
"""Per-block dense/sparse layout for M's diagonal blocks.
Blocks (connected sub-trees, each a contiguous dof range) are classified into three per-block
categories by coupling and size:
- simple: a decoupled block (M is diagonal -- a "simple body" like a point mass on orthogonal
slides) needs no factorization, just D = 1/diag, so it bypasses both factor paths.
- dense: a coupled block small enough for a dense tile-Cholesky (size <= M_BLOCK_DENSE_MAX).
- sparse: a coupled block too large for a tile, via the sparse LDL factor.
Dense block factors are packed back to back (block k's b*b factor at the prefix sum of preceding
dense block areas). Returns:
total: packed length of the dense region (also the offset of the LDL region)
dof_adr: per-dof packed offset within the dense region (0 for non-dense dofs)
blocks: all (start, size) blocks
dense_blocks: (start, size) blocks using the packed dense layout
dof_dense: per-dof flag, 1 if the dof's block is dense
dof_simple: per-dof flag, 1 if the dof's block is simple (diagonal)
has_dense / has_simple / has_sparse: whether any block falls in that category
Blocks use scalar Cholesky through six DOFs, tile Cholesky through M_BLOCK_DENSE_MAX, and sparse
LDL beyond that. Compact diagonal blocks also use the scalar path without allocating a factor.
"""
nv = mjm.nv
blocks = _m_blocks(mjm)
allow_dense = _m_allow_dense(mjm)
rownnz = mjm.M_rownnz
dof_adr = np.zeros(nv, dtype=np.int32)
dof_dense = np.zeros(nv, dtype=np.int32)
dof_simple = np.zeros(nv, dtype=np.int32)
dense_blocks = []
dof_adr = np.full(nv, types.Q_LD_BLOCK_SPARSE, dtype=np.int32)
scalar_tiles = {}
gather_tiles = {}
off = 0
has_sparse = False
for start, size in blocks:
coupled = bool(np.max(rownnz[start : start + size]) > 1)
if not coupled:
dof_simple[start : start + size] = 1
elif allow_dense and size <= types.M_BLOCK_DENSE_MAX:
dense_blocks.append((start, size))
last = start + size - 1
madr = int(mjm.M_rowadr[start])
nnz = int(mjm.M_rowadr[last] + mjm.M_rownnz[last] - madr)
compact = nnz == size
triangular = nnz == size * (size + 1) // 2
if size <= types.M_BLOCK_SCALAR_MAX and (compact or triangular):
scalar_tiles.setdefault(size, []).append(start)
if compact:
dof_adr[start : start + size] = types.Q_LD_BLOCK_COMPACT
else:
dof_adr[start : start + size] = off
off += size * size
elif size <= types.M_BLOCK_DENSE_MAX:
gather_tiles.setdefault(size, []).append(start)
dof_adr[start : start + size] = off
dof_dense[start : start + size] = 1
off += size * size
else:
has_sparse = True
for starts in scalar_tiles.values():
starts.sort(key=lambda start: dof_adr[start] >= 0)
return {
"total": off,
"dof_adr": dof_adr,
"blocks": blocks,
"dense_blocks": dense_blocks,
"dof_dense": dof_dense,
"dof_simple": dof_simple,
"has_dense": len(dense_blocks) > 0,
"has_simple": bool(dof_simple.any()),
"has_sparse": has_sparse,
"scalar_tiles": scalar_tiles,
"gather_tiles": gather_tiles,
"has_sparse": bool(np.any(dof_adr == types.Q_LD_BLOCK_SPARSE)),
}
@@ -344,9 +347,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
if (mjm.opt.enableflags & mujoco.mjtEnableBit.mjENBL_SLEEP) and (mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX).any():
raise NotImplementedError("Flex equality constraints are not supported with sleeping enabled.")
if mjm.nflex > 0 and (mjm.flex_interp < 0).any():
raise NotImplementedError("Flex interpolation order < 0 (shell/quad elements) is not supported.")
if mjm.opt.noslip_iterations > 0:
raise NotImplementedError(f"noslip solver not implemented.")
@@ -359,6 +359,19 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
if (mjm.sensor_plugin != -1).any():
raise NotImplementedError("Sensor plugins not supported.")
if mjm.nflex > 0:
for fi in range(mjm.nflex):
if abs(mjm.flex_interp[fi]) == 2:
raise NotImplementedError("Quadratic flex interpolation (dof=quadratic) is not supported.")
if mjm.flex_interp[fi] >= 0:
continue
bendingadr = mjm.flex_bendingadr[fi]
if bendingadr < 0:
continue
nedge = int(mjm.flex_bending[bendingadr])
if nedge > 0 and mjm.flex_damping[fi] > 0.0:
warnings.warn("Bending damping is not yet supported for interpolated flex shells.")
# array sizes may change in the future
if mujoco.mjNPOLY != 2:
warnings.warn(f"mujoco.mjNPOLY is {mujoco.mjNPOLY}, expected 2. Higher order polynomials may not be supported correctly.")
@@ -445,7 +458,11 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
m.callback = types.Callback()
m.nv_pad = _get_padded_sizes(
mjm.nv, 0, is_sparse(mjm), types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
mjm.nv,
0,
is_sparse(mjm),
types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE,
augment_cholesky=mjm.opt.solver == mujoco.mjtSolver.mjSOL_NEWTON and mjm.nv > 32,
)[1]
m.nacttrnbody = (mjm.actuator_trntype == mujoco.mjtTrn.mjTRN_BODY).sum()
m.nsensortaxel = mjm.mesh_vertnum[mjm.sensor_objid[mjm.sensor_type == mujoco.mjtSensor.mjSENS_TACTILE]].sum()
@@ -473,7 +490,7 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
m.block_dim.solve_search_update_cg = _nv_block
m.block_dim.solve_init_search_cg = _nv_block
if mjm.nv > 500:
m.block_dim.linesearch_iterative = 512
m.block_dim.linesearch_iterative = 256
m.is_sparse = is_sparse(mjm)
m.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0
@@ -481,12 +498,12 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
# Precompute flex_cell_map
flex_cell_map = []
if mjm.nflex > 0 and hasattr(mjm, "flex_interp"):
if mjm.nflex > 0:
for fi in range(mjm.nflex):
order = abs(int(mjm.flex_interp[fi]))
if order == 0:
continue
if hasattr(mjm, "flex_edgeequality") and mjm.flex_edgeequality[fi] == 3:
if mjm.flex_edgeequality[fi] == 3:
continue
cx, cy, cz = mjm.flex_cellnum[fi]
for ci in range(cx):
@@ -556,13 +573,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
dofid = mjm.dof_parentid[dofid]
m.body_isdofancestor = body_isdofancestor
# Upper bound on a contact's Jacobian support-pair count, to size the elliptic-cone JTCJ
# launch. Use body_isdofancestor (the full dof tree), not the mass-matrix sparsity, which the
# simple-dof optimization diagonalizes -- that undercounts the support and NaNs the solve.
support_chains = [set(np.flatnonzero(row).tolist()) for row in np.unique(body_isdofancestor[mjm.geom_bodyid], axis=0)]
max_support = max((len(ci | cj) for i, ci in enumerate(support_chains) for cj in support_chains[i:]), default=0)
m.jtcj_max_pairs = max(max_support * (max_support + 1) // 2, 1)
# precalculated geom pairs
filterparent = not (mjm.opt.disableflags & types.DisableBit.FILTERPARENT)
@@ -672,6 +682,31 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
BOX = int(mujoco.mjtGeom.mjGEOM_BOX)
MESH = int(mujoco.mjtGeom.mjGEOM_MESH)
# TODO(team): remove after implementing multicontact support for CCD pairs.
if use_multiccd:
unsupported_multiccd_pairs = []
for (g1, g2), col_type in MJ_COLLISION_TABLE.items():
if g1 == types.GeomType.BOX and g2 == types.GeomType.BOX and nativeccd_disabled:
continue
if col_type == types.CollisionType.CONVEX:
if g1 in (types.GeomType.SPHERE, types.GeomType.ELLIPSOID) or g2 in (
types.GeomType.SPHERE,
types.GeomType.ELLIPSOID,
):
continue
if (g1, g2) not in (
(types.GeomType.BOX, types.GeomType.BOX),
(types.GeomType.BOX, types.GeomType.MESH),
(types.GeomType.MESH, types.GeomType.MESH),
):
if m.geom_pair_type_count[geom_trid_index(int(g1), int(g2))] > 0:
unsupported_multiccd_pairs.append((g1.name, g2.name))
if unsupported_multiccd_pairs:
warnings.warn(
"MULTICCD is enabled, but the scene contains CCD pairs without multicontact support:"
f" {unsupported_multiccd_pairs}. At most 1 contact will be generated for these pairs."
)
has_boxbox = m.geom_pair_type_count[geom_trid_index(BOX, BOX)] > 0
has_multiccd_pairs = has_boxbox or (
use_multiccd
@@ -746,7 +781,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
m.eq_flex_adr = np.nonzero(mjm.eq_type == types.EqType.FLEX)[0]
m.eq_flexstrain_adr = np.nonzero(mjm.eq_type == types.EqType.FLEXSTRAIN)[0]
m.neq_flexstrain = m.eq_flexstrain_adr.size
# Precompute flex strain Jacobian sparsity pattern
flexstrain_J_rownnz = []
flexstrain_J_colind = []
@@ -755,6 +789,7 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
for eqstrainid, eqid in enumerate(m.eq_flexstrain_adr):
f = int(mjm.eq_obj1id[eqid])
order = int(mjm.flex_interp[f])
order_abs = abs(order)
ci = int(mjm.eq_data[eqid, 0])
cj = int(mjm.eq_data[eqid, 1])
ck = int(mjm.eq_data[eqid, 2])
@@ -763,15 +798,26 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
cy = cellnum[1]
cz = cellnum[2]
nstart = mjm.flex_nodeadr[f]
ny_g = cy * order + 1
nz_g = cz * order + 1
ny_g = cy * order_abs + 1
nz_g = cz * order_abs + 1
node_bodies = [
mjm.flex_nodebodyid[nstart + (ci * order + li) * ny_g * nz_g + (cj * order + lj) * nz_g + (ck * order + lk)]
for li in range(order + 1)
for lj in range(order + 1)
for lk in range(order + 1)
]
if order < 0:
# Shell mode: 2D bilinear quad
npc = (order_abs + 1) * (order_abs + 1)
node_bodies = []
for idx in range(npc):
gidx = support.gather_face_node_index(int(cellnum[0]), int(cy), int(cz), ci, idx, order_abs)
node_bodies.append(mjm.flex_nodebodyid[nstart + gidx])
else:
# Solid mode: 3D trilinear voxel
node_bodies = [
mjm.flex_nodebodyid[
nstart + (ci * order_abs + li) * ny_g * nz_g + (cj * order_abs + lj) * nz_g + (ck * order_abs + lk)
]
for li in range(order_abs + 1)
for lj in range(order_abs + 1)
for lk in range(order_abs + 1)
]
active_dof_mask = np.any(body_isdofancestor[node_bodies, :] != 0, axis=0)
sorted_dofs = np.nonzero(active_dof_mask)[0].tolist()
@@ -903,50 +949,47 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]])
]
# Per-block dense/sparse layout (see m_block_layout). M_tiles holds the dense blocks grouped by
# size; a model may use both paths at once (e.g. one large tree + many small free joints).
# Per-block scalar/tile/sparse layout (see m_block_layout).
_lay = m_block_layout(mjm)
dof_dense = _lay["dof_dense"]
dof_simple = _lay["dof_simple"]
m.qLD_has_dense = _lay["has_dense"]
m.qLD_has_simple = _lay["has_simple"]
m.qLD_has_sparse = _lay["has_sparse"]
m.qLD_block_total = _lay["total"] # packed dense region length / offset of the LDL region
m.qLD_block_adr = _lay["dof_adr"]
m.qLD_dof_dense = dof_dense # per-dof: 1 if the dof's block is dense (packed)
m.qLD_dof_simple = dof_simple # per-dof: 1 if the dof's block is simple (diagonal -> 1/diag)
m.qLD_simple_dofs = np.nonzero(dof_simple)[0].astype(np.int32) # the simple dof indices
tiles = {}
for start, size in _lay["dense_blocks"]:
tiles.setdefault(size, []).append(start)
m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
scalar_tiles = [
types.TileSet(
adr=wp.array(_lay["scalar_tiles"][size], dtype=int),
size=size,
)
for size in sorted(_lay["scalar_tiles"])
]
gather_tiles = [
types.TileSet(adr=wp.array(_lay["gather_tiles"][size], dtype=int), size=size) for size in sorted(_lay["gather_tiles"])
]
m.M_tiles = tuple(scalar_tiles + gather_tiles)
# qLD_updates has dof tree ordering of qLD updates for the sparse LDL factor. Only sparse-block
# dofs are included; dense blocks use the packed Cholesky and never touch the LDL region.
qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
# Group sparse LDL updates by tree depth. Block-path DOFs never touch the LDL region.
sparse_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
for k in range(mjm.nv):
# skip diagonal rows
if mjm.M_rownnz[k] == 1:
continue
dof_depth[k] = dof_depth[mjm.dof_parentid[k]] + 1
if dof_dense[k]:
continue # dense block: handled by the packed Cholesky, not the LDL factor
if _lay["dof_adr"][k] != types.Q_LD_BLOCK_SPARSE:
continue
i = mjm.dof_parentid[k]
diag_k = mjm.M_rowadr[k] + mjm.M_rownnz[k] - 1
Madr_ki = diag_k - 1
while i > -1:
qLD_updates.setdefault(dof_depth[i], []).append((i, k, Madr_ki))
sparse_updates.setdefault(dof_depth[i], []).append((i, k, Madr_ki))
i = mjm.dof_parentid[i]
Madr_ki -= 1
m.qLD_updates = tuple(wp.array(qLD_updates[i], dtype=wp.vec3i) for i in sorted(qLD_updates))
m.qLD_updates = tuple(wp.array(sparse_updates[i], dtype=wp.vec3i) for i in sorted(sparse_updates))
# Build concatenated updates for fused kernel
all_updates_flat = []
level_offsets = [0]
for level in sorted(qLD_updates):
all_updates_flat.extend(qLD_updates[level])
for level in sorted(sparse_updates):
all_updates_flat.extend(sparse_updates[level])
level_offsets.append(len(all_updates_flat))
m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
m.qLD_level_offsets = level_offsets
@@ -983,9 +1026,9 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
# dense block and flat slot (row, col), store the CSR address of M[max(i,j), min(i,j)], or nC
# (out of bounds -> read as 0) for structurally absent pairs. Laid out [block, slot] so the kernel
# reads slice (block_size^2,) at offset blk * block_size^2.
for tile in m.M_tiles:
for tile in gather_tiles:
sz = tile.size
starts = np.array(tiles[sz], dtype=np.int32) # host block starts; no device round-trip
starts = np.array(_lay["gather_tiles"][sz], dtype=np.int32)
dofs = starts[:, None] + np.arange(sz)[None, :] # (nblock, sz) global dof per block row
gi = dofs[:, :, None] # (nblock, sz, 1)
gj = dofs[:, None, :] # (nblock, 1, sz)
@@ -1144,8 +1187,109 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
m.flex_vertflexid = flex_vertflexid
m.flex_shelladr = flex_shelladr
# place m on device
flex_bend_interp_map = []
flex_face_map = []
flex_faceadr = np.zeros(mjm.nflex, dtype=np.int32)
if mjm.nflex > 0:
face_offset = 0
for fi in range(mjm.nflex):
flex_faceadr[fi] = face_offset
order = mjm.flex_interp[fi]
if order >= 0:
continue
bendingadr = mjm.flex_bendingadr[fi]
if bendingadr >= 0:
nedge = int(mjm.flex_bending[bendingadr])
for e in range(nedge):
flex_bend_interp_map.append((fi, e))
cx, cy, cz = mjm.flex_cellnum[fi]
nfaces = 2 * (cy * cz + cx * cz + cx * cy)
for face_idx in range(nfaces):
flex_face_map.append((fi, face_idx))
face_offset += nfaces
if not flex_bend_interp_map:
m.nflexbend_interp = 0
m.flex_bend_interp_map = np.zeros((0, 2), dtype=np.int32)
else:
m.nflexbend_interp = len(flex_bend_interp_map)
m.flex_bend_interp_map = np.array(flex_bend_interp_map, dtype=np.int32)
if not flex_face_map:
m.nflexface = 0
m.flex_face_map = np.zeros((0, 2), dtype=np.int32)
else:
m.nflexface = len(flex_face_map)
m.flex_face_map = np.array(flex_face_map, dtype=np.int32)
m.flex_faceadr = flex_faceadr
if m.nflexface > 0:
flex_face = np.zeros((m.nflexface, 9), dtype=np.int32)
for face_id, (fi, face_elem_idx) in enumerate(flex_face_map):
order = mjm.flex_interp[fi]
order_abs = -order
cx, cy, cz = mjm.flex_cellnum[fi]
nstart = mjm.flex_nodeadr[fi]
npc = (order_abs + 1) * (order_abs + 1)
for local_idx in range(9):
if local_idx < npc:
gidx = support.gather_face_node_index(int(cx), int(cy), int(cz), int(face_elem_idx), int(local_idx), int(order_abs))
flex_face[face_id, local_idx] = nstart + gidx
else:
flex_face[face_id, local_idx] = -1
m.flex_face = flex_face
else:
m.flex_face = np.zeros((0, 9), dtype=np.int32)
sizes = {f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int}
sizes.update(
{
"nbody_branches": len(m.body_branches),
"nbranch_start": len(m.body_branch_start),
"nbody_fluid_ellipsoid": len(m.body_fluid_ellipsoid_adr),
"nbody_fluid_box": len(m.body_fluid_box_adr),
"njnt_limited_slide_hinge": len(m.jnt_limited_slide_hinge_adr),
"njnt_limited_ball": len(m.jnt_limited_ball_adr),
"ndof_tri": len(m.dof_tri_row),
"nnxn_geom_pair": len(m.nxn_geom_pair),
"nnxn_geom_pair_filtered": len(m.nxn_geom_pair_filtered),
"neq_connect": len(m.eq_connect_adr),
"neq_wld": len(m.eq_wld_adr),
"neq_jnt": len(m.eq_jnt_adr),
"neq_ten": len(m.eq_ten_adr),
"neq_flex": len(m.eq_flex_adr),
"ntendon_jnt": len(m.tendon_jnt_adr),
"ntendon_site_pair": len(m.tendon_site_pair_adr),
"ntendon_geom": len(m.tendon_geom_adr),
"ntendon_limited": len(m.tendon_limited_adr),
"nten_wrapadr_site": len(m.ten_wrapadr_site),
"nwrap_jnt": len(m.wrap_jnt_adr),
"nwrap_site": len(m.wrap_site_adr),
"nwrap_site_pair": len(m.wrap_site_pair_adr),
"nwrap_geom": len(m.wrap_geom_adr),
"nsensor_pos": len(m.sensor_pos_adr),
"nsensor_limitpos": len(m.sensor_limitpos_adr),
"nsensor_vel": len(m.sensor_vel_adr),
"nsensor_limitvel": len(m.sensor_limitvel_adr),
"nsensor_acc": len(m.sensor_acc_adr),
"nsensor_touch": len(m.sensor_touch_adr),
"nsensor_limitfrc": len(m.sensor_limitfrc_adr),
"nsensor_tendonactfrc": len(m.sensor_tendonactfrc_adr),
"nsensor_collision_start_adr": len(m.sensor_collision_start_adr),
"nqLD_all_updates": len(m.qLD_all_updates),
"nqLD_level_offsets": len(m.qLD_level_offsets),
"nM_fullm": len(m.M_fullm_i),
"nM_fullm_upper": len(m.M_fullm_upper_i),
"nqD_fullm": len(m.qD_fullm_i),
"nv_plus_1": len(m.M_mulm_rowadr),
"nM_mulm": len(m.M_mulm_col),
"nflexelem_geom_pair_filtered": len(m.flexelem_geom_pair_filtered),
"nflexvert_geom_pair_filtered": len(m.flexvert_geom_pair_filtered),
}
)
for f in dataclasses.fields(types.Model):
if _is_array_spec(f.type):
batch_size = batch_sizes.get(f.name, 1)
@@ -1155,24 +1299,20 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
return m
def _get_padded_sizes(nv: int, njmax: int, is_sparse: bool, tile_size: int):
# if dense - we just pad to the next multiple of 4 for nv, to get the fast load path.
# we pad to the next multiple of tile_size for njmax to avoid out of bounds accesses.
# if sparse - we pad to the next multiple of tile_size for njmax, and nv.
def _get_padded_sizes(nv: int, njmax: int, is_sparse: bool, tile_size: int, augment_cholesky: bool = False):
def round_up(x, multiple):
return ((x + multiple - 1) // multiple) * multiple
njmax_padded = round_up(njmax, tile_size)
nv_padded = round_up(nv, tile_size) if (is_sparse or nv > 32) else round_up(nv, 4)
nv_padded = round_up(nv + int(augment_cholesky), tile_size) if (is_sparse or nv > 32) else round_up(nv, 4)
return njmax_padded, nv_padded
def _nvmax_pad(nvmax: int) -> int:
"""Round nvmax up to the dense tile size so the blocked Cholesky never overruns its tile."""
"""Reserve an augmented column and round nvmax up to the dense tile size."""
t = types.TILE_SIZE_JTDAJ_DENSE
return ((max(nvmax, 1) + t - 1) // t) * t
return ((max(nvmax, 1) + t) // t) * t
def _default_nconmax(mjm: mujoco.MjModel, mjd: Optional[mujoco.MjData] = None) -> int:
@@ -1218,6 +1358,18 @@ def _body_pair_nnz(mjm: mujoco.MjModel, body1: int, body2: int) -> int:
return nnz
def _body_set_nnz(mjm: mujoco.MjModel, bodies) -> int:
"""Returns the number of unique DOFs in the kinematic tree union of a set of bodies."""
active_dofs = set()
for b in bodies:
b = mjm.body_weldid[b]
da = mjm.body_dofadr[b] + mjm.body_dofnum[b] - 1
while da >= 0:
active_dofs.add(da)
da = mjm.dof_parentid[da]
return len(active_dofs)
def _default_njmax_nnz(mjm: mujoco.MjModel, nconmax: int, njmax: int) -> int:
"""Returns a heuristic estimate for the number of non-zeros in the sparse constraint Jacobian.
@@ -1283,7 +1435,7 @@ def _default_njmax_nnz(mjm: mujoco.MjModel, nconmax: int, njmax: int) -> int:
elif eq_type == mujoco.mjtEq.mjEQ_FLEXSTRAIN:
# strain constraints: each cell produces neig rows, each dense (nv)
obj1id = mjm.eq_obj1id[i]
if obj1id < mjm.nflex and hasattr(mjm, "flex_stiffnessadr"):
if obj1id < mjm.nflex:
# estimate neig from stiffness data
adr = mjm.flex_stiffnessadr[obj1id]
neig = int(mjm.flex_stiffness[adr])
@@ -1351,10 +1503,57 @@ def _default_njmax_nnz(mjm: mujoco.MjModel, nconmax: int, njmax: int) -> int:
if (fct & ca) or (ct & fca):
geom_bodies.add(mjm.geom_bodyid[g])
for fb in flex_bodies:
if mjm.flex_interp[fi] == 0:
for fb in flex_bodies:
for gb in geom_bodies:
if fb != gb:
max_contact_nnz = max(max_contact_nnz, _body_pair_nnz(mjm, fb, gb))
else:
order = abs(mjm.flex_interp[fi])
is_shell = mjm.flex_interp[fi] < 0
cx, cy, cz = mjm.flex_cellnum[fi]
nstart = mjm.flex_nodeadr[fi]
dim = mjm.flex_dim[fi]
nx = cx * order + 1
ny = cy * order + 1 if dim > 1 else 1
nz = cz * order + 1 if dim > 2 else 1
ci, cj, ck = cx // 2, cy // 2, cz // 2
cell_bodies = set()
for li in range(order + 1):
for lj in range(order + 1 if dim > 1 else 1):
for lk in range(order + 1 if dim > 2 else 1):
gi = ci + li
gj = cj + lj
gk = ck + lk
is_interior = False
if is_shell:
is_interior = (
(gi > 0 and gi < cx * order)
and (gj > 0 and gj < cy * order if dim > 1 else True)
and (gk > 0 and gk < cz * order if dim > 2 else True)
)
if is_interior:
for bi in (0, gi, nx - 1):
for bj in (0, gj, ny - 1 if dim > 1 else 0):
for bk in (0, gk, nz - 1 if dim > 2 else 0):
if (
bi == 0
or bi == nx - 1
or (dim > 1 and (bj == 0 or bj == ny - 1))
or (dim > 2 and (bk == 0 or bk == nz - 1))
):
node_idx = bi * ny * nz + bj * nz + bk
cell_bodies.add(mjm.flex_nodebodyid[nstart + node_idx])
else:
node_idx = gi * ny * nz + gj * nz + gk
cell_bodies.add(mjm.flex_nodebodyid[nstart + node_idx])
for gb in geom_bodies:
if fb != gb:
max_contact_nnz = max(max_contact_nnz, _body_pair_nnz(mjm, fb, gb))
max_contact_nnz = max(max_contact_nnz, _body_set_nnz(mjm, cell_bodies | {gb}))
# flex self-collision
if mjm.flex_selfcollide[fi]:
@@ -1455,11 +1654,13 @@ def _allocate_compact_arrays(
d.cdof_tri_row = wp.empty(0, dtype=int)
d.cdof_tri_col = wp.empty(0, dtype=int)
alloc_cJ = compact and not is_sparse(mjm)
d.cM = wp.empty((nw, nvp, nvp), dtype=float)
d.cqLD = wp.empty((nw, nvp, nvp), dtype=float)
d.crhs = wp.empty((nw, nvp, 1), dtype=float)
d.cx = wp.empty((nw, nvp, 1), dtype=float)
d.cJ = wp.empty((nw, njp, nvp), dtype=float)
d.cJ = wp.empty((nw, njp, nvp), dtype=float) if alloc_cJ else wp.empty((0, 0, 0), dtype=float)
d.cMa = wp.empty((nw, nvp), dtype=float)
d.cqfrc_smooth = wp.empty((nw, nvp), dtype=float)
d.cqacc_smooth = wp.empty((nw, nvp), dtype=float)
@@ -1569,13 +1770,28 @@ def make_data(
sizes["nmaxcondim"] = np.concatenate(condim_arrays).max()
sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1))
tile_size = types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, is_sparse(mjm), tile_size)
sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(
mjm.nv,
njmax,
is_sparse(mjm),
tile_size,
augment_cholesky=mjm.opt.solver == mujoco.mjtSolver.mjSOL_NEWTON and mjm.nv > 32,
)
sizes["nworld"] = nworld
sizes["naconmax"] = naconmax
sizes["njmax"] = njmax
sizes["nvmax"] = nvmax
sizes["nvmax_pad"] = _nvmax_pad(nvmax)
sizes["nvmax_pad_sq"] = sizes["nvmax_pad"] * sizes["nvmax_pad"]
sizes["nflexintcell"] = _get_nflexintcell(mjm)
sizes["nflexface"] = _get_nflexface(mjm)
# qLD holds the factor: a packed dense region for dense blocks followed
# by an nC-length LDL region for sparse blocks (present only when some block is sparse). Either
# region may be empty (pure dense / pure sparse).
_lay = m_block_layout(mjm)
qld_total = _lay["total"] + (mjm.nC if _lay["has_sparse"] else 0)
sizes["qld_total"] = qld_total
if njmax_nnz is None:
if is_sparse(mjm):
@@ -1592,7 +1808,7 @@ def make_data(
contact = types.Contact(**contact_kwargs)
contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes)
efc = _create_constraint(mjm, nworld, njmax, sizes)
if is_sparse(mjm):
efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
@@ -1627,8 +1843,6 @@ def make_data(
"nvmax_pad": sizes["nvmax_pad"],
"njmax_pad": sizes["njmax_pad"],
"njmax_nnz": njmax_nnz,
"M": None,
"qLD": None,
# world body
"xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
"xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
@@ -1672,14 +1886,6 @@ def make_data(
d = types.Data(**d_kwargs)
# qLD holds the factor: a packed dense region for dense blocks followed
# by an nC-length LDL region for sparse blocks (present only when some block is sparse). Either
# region may be empty (pure dense / pure sparse).
d.M = wp.zeros((nworld, mjm.nC), dtype=float)
_lay = m_block_layout(mjm)
qld_total = _lay["total"] + (mjm.nC if _lay["has_sparse"] else 0)
d.qLD = wp.zeros((nworld, qld_total), dtype=float)
_allocate_island_arrays(mjm, d, nworld, njmax, island_alloc, mjd)
_allocate_compact_arrays(mjm, d, nworld, sizes["nvmax_pad"], sizes["njmax_pad"], compact_alloc)
d.ncdof.zero_()
@@ -1789,12 +1995,20 @@ def put_data(
sizes["nmaxcondim"] = np.concatenate(condim_arrays).max()
sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1))
tile_size = types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(mjm.nv, njmax, is_sparse(mjm), tile_size)
sizes["njmax_pad"], sizes["nv_pad"] = _get_padded_sizes(
mjm.nv,
njmax,
is_sparse(mjm),
tile_size,
augment_cholesky=mjm.opt.solver == mujoco.mjtSolver.mjSOL_NEWTON and mjm.nv > 32,
)
sizes["nworld"] = nworld
sizes["naconmax"] = naconmax
sizes["njmax"] = njmax
sizes["nvmax"] = nvmax
sizes["nvmax_pad"] = _nvmax_pad(nvmax)
sizes["nvmax_pad_sq"] = sizes["nvmax_pad"] * sizes["nvmax_pad"]
sizes["nflexface"] = _get_nflexface(mjm)
if njmax_nnz is None:
if is_sparse(mjm):
@@ -1844,9 +2058,7 @@ def put_data(
contact.geomcollisionid = wp.empty((naconmax,), dtype=int) # TODO(team): set values
# create efc
efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, mjd)
efc = _create_constraint(mjm, nworld, njmax, sizes, mjd)
# make_constraint builds the block list in-kernel; put_data does not run it, so build it here
# -- otherwise solving a put_data state would assemble an empty J^T D J.
@@ -1912,7 +2124,6 @@ def put_data(
"njmax_nnz": njmax_nnz,
# fields set after initialization:
"solver_niter": None,
"M": None,
"qLD": None,
"nacon": None,
# island arrays
@@ -1946,21 +2157,23 @@ def put_data(
d = types.Data(**d_kwargs)
d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
d.M = wp.array(np.full((nworld, mjm.nC), mjd.M), dtype=float)
# qLD = [packed dense-block Cholesky | nC LDL region]. Dense blocks store their upper Cholesky
# qLD = [packed block Cholesky | nC LDL region]. Block factors store their upper Cholesky
# packed; the LDL region (present iff some block is sparse) holds MuJoCo's full L'DL factor (only
# its sparse-block entries are read by the solve).
lay = m_block_layout(mjm)
qld_total = lay["total"] + (mjm.nC if lay["has_sparse"] else 0)
qLD = np.zeros(qld_total, dtype=np.float32)
if lay["has_dense"]:
if lay["total"]:
Mfull = np.zeros((mjm.nv, mjm.nv))
mujoco.mju_sym2dense(Mfull, mjd.M, mjm.M_rownnz, mjm.M_rowadr, mjm.M_colind)
for start, size in lay["dense_blocks"]:
off = lay["dof_adr"][start]
blk = Mfull[start : start + size, start : start + size]
if blk.any():
qLD[off : off + size * size] = np.linalg.cholesky(blk).T.reshape(-1)
for size, starts in list(lay["scalar_tiles"].items()) + list(lay["gather_tiles"].items()):
for start in starts:
off = lay["dof_adr"][start]
if off < 0:
continue
blk = Mfull[start : start + size, start : start + size]
if blk.any():
qLD[off : off + size * size] = np.linalg.cholesky(blk).T.reshape(-1)
if lay["has_sparse"]:
qLD[lay["total"] :] = mjd.qLD
d.qLD = wp.array(np.full((nworld, qld_total), qLD), dtype=float)
@@ -2129,9 +2342,8 @@ def get_data_into(
result.M[:] = d.M.numpy()[world_id]
_lay = m_block_layout(mjm)
if _lay["has_dense"] or _lay["has_simple"]:
# d.qLD is not MuJoCo's LDL: dense blocks are a packed Cholesky and simple blocks are factored
# into qLDiagInv (their LDL slots are never written). Recompute the LDL factor from M.
if _lay["scalar_tiles"] or _lay["gather_tiles"]:
# Block factors do not use MuJoCo's LDL representation.
mujoco.mj_factorM(mjm, result)
else:
# Pure sparse: qLD is exactly MuJoCo's nC LDL factor.
@@ -2227,7 +2439,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
"""
sleep_enabled = bool(m.opt.enableflags & types.EnableBit.SLEEP)
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def reset_xfrc_applied(reset_in: wp.array[bool], xfrc_applied_out: wp.array2d[wp.spatial_vector]):
worldid, bodyid, elemid = wp.tid()
@@ -2237,7 +2449,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
xfrc_applied_out[worldid, bodyid][elemid] = 0.0
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def reset_M(reset_in: wp.array[bool], M_out: wp.array2d[float]):
worldid, elemid = wp.tid()
@@ -2247,7 +2459,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
M_out[worldid, elemid] = 0.0
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def reset_nworld(
# Model:
nq: int,
@@ -2329,7 +2541,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
userdata_out[worldid, i] = 0.0
overflow_out[worldid] = 0
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def reset_mocap(
# Model:
body_mocapid: wp.array[int],
@@ -2353,7 +2565,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
mocap_pos_out[worldid, mocapid] = body_pos[worldid % body_pos.shape[0], bodyid]
mocap_quat_out[worldid, mocapid] = body_quat[worldid % body_quat.shape[0], bodyid]
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def reset_contact(
# Data in:
nacon_in: wp.array[int],
@@ -2412,7 +2624,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
contact_type_out[conid] = 0
contact_geomcollisionid_out[conid] = 0
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def reset_sleep(
# Model:
nv: int,
@@ -3051,14 +3263,14 @@ def _compute_dof_M0(
@wp.kernel
def _resolve_dampratio(
actuator_biastype: wp.array[int],
actuator_gainprm: wp.array2d[types.vec10f],
actuator_gainprm: wp.array2d[types.vec10],
moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d[int],
moment_colind_in: wp.array2d[int],
actuator_moment_in: wp.array2d[float],
dof_M0_in: wp.array2d[float],
nv: int,
actuator_biasprm: wp.array2d[types.vec10f],
actuator_biasprm: wp.array2d[types.vec10],
):
worldid, actid = wp.tid()
biastype = actuator_biastype[actid]
@@ -3149,12 +3361,13 @@ def set_const_fixed(m: types.Model, d: types.Data):
m: The model containing kinematic and dynamic information (device).
d: The data object containing the current state and output arrays (device).
"""
wp.launch(_init_subtreemass, dim=(d.nworld, m.nbody), inputs=[m.body_mass], outputs=[m.body_subtreemass])
nworld_subtreemass = m.body_subtreemass.shape[0]
wp.launch(_init_subtreemass, dim=(nworld_subtreemass, m.nbody), inputs=[m.body_mass], outputs=[m.body_subtreemass])
for i in reversed(range(len(m.body_tree))):
body_tree = m.body_tree[i]
wp.launch(
_accumulate_subtreemass,
dim=(d.nworld, body_tree.size),
dim=(nworld_subtreemass, body_tree.size),
inputs=[m.body_parentid, m.body_subtreemass, body_tree],
)
@@ -3198,16 +3411,16 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
# Compute meaninertia from M diagonal at qpos0
wp.launch(
_compute_meaninertia,
dim=d.nworld,
dim=m.stat.meaninertia.shape[0],
inputs=[m.nv, m.M_rownnz, m.M_rowadr, d.M],
outputs=[m.stat.meaninertia],
)
wp.launch(_copy_tendon_length0, dim=(d.nworld, m.ntendon), inputs=[d.ten_length], outputs=[m.tendon_length0])
wp.launch(_copy_tendon_length0, dim=(m.tendon_length0.shape[0], m.ntendon), inputs=[d.ten_length], outputs=[m.tendon_length0])
wp.launch(
_compute_eq_data0,
dim=(d.nworld, m.neq),
dim=(m.eq_data.shape[0], m.neq),
inputs=[m.eq_type, m.eq_obj1id, m.eq_obj2id, m.eq_objtype, d.xpos, d.xquat, d.xmat],
outputs=[m.eq_data],
)
@@ -3229,7 +3442,7 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
wp.launch(
_finalize_dof_invweight0,
dim=(d.nworld, m.nv),
dim=(m.dof_invweight0.shape[0], m.nv),
inputs=[m.dof_jntid, m.jnt_type, m.jnt_dofadr, dof_A_diag],
outputs=[m.dof_invweight0],
)
@@ -3272,7 +3485,7 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
wp.launch(
_finalize_body_invweight0,
dim=(d.nworld, m.nbody),
dim=(m.body_invweight0.shape[0], m.nbody),
inputs=[m.body_weldid, body_A_diag],
outputs=[m.body_invweight0],
)
@@ -3295,21 +3508,23 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
smooth.solve_m(m, d, ten_result_vec, ten_J_vec)
wp.launch(
_compute_tendon_dot_product,
dim=d.nworld,
dim=m.tendon_invweight0.shape[0],
inputs=[m.ten_J_rownnz, m.ten_J_rowadr, m.ten_J_colind, tenid, d.ten_J, ten_result_vec],
outputs=[m.tendon_invweight0],
)
nworld_cam = np.max([m.cam_pos0.shape[0], m.cam_poscom0.shape[0], m.cam_mat0.shape[0]])
wp.launch(
_compute_cam_pos0,
dim=(d.nworld, m.ncam),
dim=(nworld_cam, m.ncam),
inputs=[m.cam_bodyid, m.cam_targetbodyid, d.cam_xpos, d.cam_xmat, d.xpos, d.subtree_com],
outputs=[m.cam_pos0, m.cam_poscom0, m.cam_mat0],
)
nworld_light = np.max([m.light_pos0.shape[0], m.light_poscom0.shape[0], m.light_dir0.shape[0]])
wp.launch(
_compute_light_pos0,
dim=(d.nworld, m.nlight),
dim=(nworld_light, m.nlight),
inputs=[m.light_bodyid, m.light_targetbodyid, d.light_xpos, d.light_xdir, d.xpos, d.subtree_com],
outputs=[m.light_pos0, m.light_poscom0, m.light_dir0],
)
@@ -3327,7 +3542,9 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
outputs=[act_moment_vec],
)
smooth.solve_m(m, d, act_result_vec, act_moment_vec)
wp.launch(_compute_actuator_acc0, dim=d.nworld, inputs=[actid, m.nv, act_result_vec], outputs=[m.actuator_acc0])
wp.launch(
_compute_actuator_acc0, dim=m.actuator_acc0.shape[0], inputs=[actid, m.nv, act_result_vec], outputs=[m.actuator_acc0]
)
# resolve dampratio: compute dof_M0, then convert dampratio to damping
if m.nu > 0 and m.nv > 0:
@@ -3340,7 +3557,7 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
)
wp.launch(
_resolve_dampratio,
dim=(d.nworld, m.nu),
dim=(m.actuator_biasprm.shape[0], m.nu),
inputs=[
m.actuator_biastype,
m.actuator_gainprm,
@@ -3388,7 +3605,7 @@ def set_const_spring(m: types.Model, d: types.Data, restore: bool = True):
wp.launch(
_resolve_tendon_lengthspring,
dim=(d.nworld, m.ntendon),
dim=(m.tendon_lengthspring.shape[0], m.ntendon),
inputs=[d.ten_length],
outputs=[m.tendon_lengthspring],
)
@@ -3701,6 +3918,7 @@ def create_render_context(
render_depth: list[bool] | bool | None = None,
render_seg: list[bool] | bool | None = None,
use_textures: bool = True,
use_fast_math: bool = True,
use_shadows: bool = False,
use_ambient_lighting: bool = True,
enabled_geom_groups: list[int] = [0, 1, 2],
@@ -3726,6 +3944,7 @@ def create_render_context(
render_seg: Whether to render segmentation (per-pixel object ID/type pairs).
If None, uses the MuJoCo model values.
use_textures: Whether to use textures.
use_fast_math: Whether to enable fast math for the render kernel.
use_shadows: Whether to use shadows.
use_ambient_lighting: Top-level ambient switch. When False, skips all
ambient contributions, including headlight ambient,
@@ -3948,6 +4167,7 @@ def create_render_context(
cam_res=cam_res_arr,
cam_id_map=wp.array(active_cam_indices, dtype=int),
use_textures=use_textures,
use_fast_math=use_fast_math,
use_shadows=use_shadows,
use_ambient_lighting=use_ambient_lighting,
background_color=render_util.pack_rgba_to_uint32(
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+2
View File
@@ -22,6 +22,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
from mujoco.mjx.third_party.mujoco_warp._src.types import OverflowType
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"default_grid_stride": False})
@wp.kernel
def _tree_edges(
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+343 -64
View File
@@ -24,6 +24,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
@@ -593,7 +594,6 @@ def _flex_elasticity(
# Model:
nflex: int,
opt_timestep: wp.array[float],
body_dofadr: wp.array[int],
flex_dim: wp.array[int],
flex_vertadr: wp.array[int],
flex_edgeadr: wp.array[int],
@@ -609,13 +609,14 @@ def _flex_elasticity(
flex_stiffness: wp.array[float],
flex_damping: wp.array[float],
# Data in:
xipos_in: wp.array2d[wp.vec3],
flexvert_xpos_in: wp.array2d[wp.vec3],
flexedge_length_in: wp.array2d[float],
flexedge_velocity_in: wp.array2d[float],
# In:
dsbl_damper: bool,
# Data out:
qfrc_spring_out: wp.array2d[float],
# Out:
flex_spring_body_force_out: wp.array2d[wp.spatial_vector],
):
worldid, elemid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -677,7 +678,7 @@ def _flex_elasticity(
elongation[e] = deformed * deformed - reference * reference + (deformed * deformed - previous * previous) * kD
metric = wp.matrix(0.0, shape=(6, 6))
stiffness_size = nedge * (nedge + 1) / 2
stiffness_size = 21
stiffness_adr = stiffness_adr_base + local_elemid * stiffness_size
id = int(0)
for ed1 in range(nedge):
@@ -696,15 +697,21 @@ def _flex_elasticity(
for v in range(nvert):
vert = flex_elem[elem_data_adr + v]
bodyid = flex_vertbodyid[flex_vertadr[f] + vert]
for x in range(3):
wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[v, x])
frc = force[v]
node_pos = flexvert_xpos_in[worldid, flex_vertadr[f] + vert]
body_xipos = xipos_in[worldid, bodyid]
offset = body_xipos - node_pos
spatial_frc = wp.spatial_vector(frc, -wp.cross(offset, frc))
wp.atomic_add(flex_spring_body_force_out, worldid, bodyid, spatial_frc)
@wp.kernel
def _flex_bending(
# Model:
nflex: int,
body_dofadr: wp.array[int],
body_rootid: wp.array[int],
flex_dim: wp.array[int],
flex_vertadr: wp.array[int],
flex_edgeadr: wp.array[int],
@@ -714,17 +721,23 @@ def _flex_bending(
flex_edge: wp.array[wp.vec2i],
flex_edgeflap: wp.array[wp.vec2i],
flex_bending: wp.array[float],
flex_damping: wp.array[float],
# Data in:
xipos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d[wp.vec3],
flexvert_xpos_in: wp.array2d[wp.vec3],
# Data out:
qfrc_spring_out: wp.array2d[float],
cvel_in: wp.array2d[wp.spatial_vector],
# In:
dsbl_damper: bool,
# Out:
flex_spring_body_force_out: wp.array2d[wp.spatial_vector],
flex_damper_body_force_out: wp.array2d[wp.spatial_vector],
):
worldid, edgeid = wp.tid()
nvert = 4
for i in range(nflex):
locid = edgeid - flex_edgeadr[i]
if locid >= 0 and locid < flex_edgenum[i]:
eid = edgeid - flex_edgeadr[i]
if eid >= 0 and eid < flex_edgenum[i]:
f = i
break
@@ -745,29 +758,64 @@ def _flex_bending(
flex_vertadr[f] + flex_edgeflap[edgeid][1],
)
frc = wp.matrix(0.0, shape=(4, 3))
if flex_bending[bendingadr + 16]:
frc = mat43()
if flex_bending[bendingadr + 17 * eid + 16]:
v0 = flexvert_xpos_in[worldid, v[0]]
v1 = flexvert_xpos_in[worldid, v[1]]
v2 = flexvert_xpos_in[worldid, v[2]]
v3 = flexvert_xpos_in[worldid, v[3]]
frc[1] = wp.cross(v2 - v0, v3 - v0)
frc[2] = wp.cross(v3 - v0, v1 - v0)
frc[3] = wp.cross(v1 - v0, v2 - v0)
ed0 = v1 - v0
ed1 = v2 - v0
ed2 = v3 - v0
frc[1] = wp.cross(ed1, ed2)
frc[2] = wp.cross(ed2, ed0)
frc[3] = wp.cross(ed0, ed1)
frc[0] = -(frc[1] + frc[2] + frc[3])
force = wp.matrix(0.0, shape=(nvert, 3))
for i in range(nvert):
for x in range(3):
acc = float(0.0)
for j in range(nvert):
acc += flex_bending[bendingadr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] = -(acc + flex_bending[bendingadr + 16] * frc[i, x])
# Gather velocities if damping is enabled
vel = mat43()
if not dsbl_damper and flex_damping[f] > 0.0:
for j in range(4):
bodyid_j = flex_vertbodyid[v[j]]
cvel_j = cvel_in[worldid, bodyid_j]
omega_j = wp.spatial_top(cvel_j)
vcom_j = wp.spatial_bottom(cvel_j)
com_j = subtree_com_in[worldid, body_rootid[bodyid_j]]
r_j = flexvert_xpos_in[worldid, v[j]] - com_j
vel[j] = vcom_j + wp.cross(omega_j, r_j)
for i in range(nvert):
bodyid = flex_vertbodyid[v[i]]
force_spring = mat43()
force_damper = mat43()
for i in range(4):
for x in range(3):
wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[i, x])
acc_spring = float(0.0)
acc_damper = float(0.0)
for j in range(4):
coeff = flex_bending[bendingadr + 17 * eid + 4 * i + j]
acc_spring += coeff * flexvert_xpos_in[worldid, v[j]][x]
if not dsbl_damper and flex_damping[f] > 0.0:
acc_damper += coeff * vel[j, x]
force_spring[i, x] = -(acc_spring + flex_bending[bendingadr + 17 * eid + 16] * frc[i, x])
if not dsbl_damper and flex_damping[f] > 0.0:
force_damper[i, x] = -acc_damper
for i in range(4):
bodyid = flex_vertbodyid[v[i]]
frc_s = force_spring[i]
node_pos = flexvert_xpos_in[worldid, v[i]]
body_xipos = xipos_in[worldid, bodyid]
offset = body_xipos - node_pos
spatial_frc_s = wp.spatial_vector(frc_s, -wp.cross(offset, frc_s))
wp.atomic_add(flex_spring_body_force_out, worldid, bodyid, spatial_frc_s)
if not dsbl_damper and flex_damping[f] > 0.0:
frc_d = force_damper[i] * flex_damping[f]
spatial_frc_d = wp.spatial_vector(frc_d, -wp.cross(offset, frc_d))
wp.atomic_add(flex_damper_body_force_out, worldid, bodyid, spatial_frc_d)
@wp.kernel
@@ -775,8 +823,6 @@ def _flex_passive_interp(
# Model:
nflex: int,
body_rootid: wp.array[int],
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
flex_interp: wp.array[int],
flex_cellnum: wp.array[wp.vec3i],
flex_nodeadr: wp.array[int],
@@ -790,16 +836,16 @@ def _flex_passive_interp(
flex_centered: wp.array[bool],
flex_cell_map: wp.array[wp.vec4i],
# Data in:
xipos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d[wp.vec3],
cvel_in: wp.array2d[wp.spatial_vector],
flexnode_xpos_in: wp.array2d[wp.vec3],
# In:
dsbl_spring: bool,
dsbl_damper: bool,
# Data out:
qfrc_spring_out: wp.array2d[float],
qfrc_damper_out: wp.array2d[float],
# Out:
flex_spring_body_force_out: wp.array2d[wp.spatial_vector],
flex_damper_body_force_out: wp.array2d[wp.spatial_vector],
displ_scratch_out: wp.array3d[wp.vec3],
vel_corot_scratch_out: wp.array3d[wp.vec3],
):
@@ -915,35 +961,223 @@ def _flex_passive_interp(
# Scale damper force by damping coefficient
frc_damper_world = frc_damper_world * flex_damping[f]
# Apply forces to body DOFs (fast path: nodes at body origin)
dofnum_i = body_dofnum[bodyid_i]
dofadr_i = body_dofadr[bodyid_i]
if dofnum_i > 0:
centered = flex_centered[f]
node_local = flex_node[nstart + gidx_i]
at_origin = node_local[0] == 0.0 and node_local[1] == 0.0 and node_local[2] == 0.0
# Apply forces to body
node_pos = flexnode_xpos_in[worldid, nstart + gidx_i]
body_xipos = xipos_in[worldid, bodyid_i]
if centered or at_origin:
for x in range(3):
if x < dofnum_i:
if not dsbl_spring:
wp.atomic_add(
qfrc_spring_out,
worldid,
dofadr_i + x,
frc_spring_world[x],
)
if not dsbl_damper:
wp.atomic_add(
qfrc_damper_out,
worldid,
dofadr_i + x,
frc_damper_world[x],
)
offset = body_xipos - node_pos
if not dsbl_spring:
spatial_frc_s = wp.spatial_vector(frc_spring_world, -wp.cross(offset, frc_spring_world))
wp.atomic_add(flex_spring_body_force_out, worldid, bodyid_i, spatial_frc_s)
if not dsbl_damper:
spatial_frc_d = wp.spatial_vector(frc_damper_world, -wp.cross(offset, frc_damper_world))
wp.atomic_add(flex_damper_body_force_out, worldid, bodyid_i, spatial_frc_d)
idx_i += 1
@wp.func
def _apply_face_forces(
# Model:
flex_nodebodyid: wp.array[int],
flex_face: wp.array2d[int],
# Data in:
xipos_in: wp.array2d[wp.vec3],
flexnode_xpos_in: wp.array2d[wp.vec3],
# In:
face_id: int,
local_coords: wp.vec2,
wt1: wp.vec3,
wt2: wp.vec3,
stiffness_scale: float,
order_abs: int,
worldid: int,
# Out:
body_force_out: wp.array2d[wp.spatial_vector],
):
idx = int(0)
for l0 in range(3):
if l0 > order_abs:
continue
for l1 in range(3):
if l1 > order_abs:
continue
g0 = support.dphi2D(local_coords[0], l0, local_coords[1], l1, order_abs, 0)
g1 = support.dphi2D(local_coords[0], l0, local_coords[1], l1, order_abs, 1)
gidx = flex_face[face_id, idx]
frc = (wt2 * g0 - wt1 * g1) * stiffness_scale
bid = flex_nodebodyid[gidx]
node_pos = flexnode_xpos_in[worldid, gidx]
body_xipos = xipos_in[worldid, bid]
offset = body_xipos - node_pos
spatial_frc = wp.spatial_vector(frc, -wp.cross(offset, frc))
wp.atomic_add(body_force_out, worldid, bid, spatial_frc)
idx += 1
@wp.kernel
def _flex_passive_bend_interp(
# Model:
nflex: int,
flex_interp: wp.array[int],
flex_cellnum: wp.array[wp.vec3i],
flex_nodeadr: wp.array[int],
flex_nodenum: wp.array[int],
flex_bendingadr: wp.array[int],
flex_nodebodyid: wp.array[int],
flex_node: wp.array[wp.vec3],
flex_bending: wp.array[float],
flex_centered: wp.array[bool],
flex_faceadr: wp.array[int],
flex_bend_interp_map: wp.array[wp.vec2i],
flex_face: wp.array2d[int],
# Data in:
xipos_in: wp.array2d[wp.vec3],
flexnode_xpos_in: wp.array2d[wp.vec3],
face_xpos_in: wp.array3d[wp.vec3],
face_quat_in: wp.array2d[wp.quat],
# Out:
flex_spring_body_force_out: wp.array2d[wp.spatial_vector],
):
worldid, bend_edge_id = wp.tid()
mapping = flex_bend_interp_map[bend_edge_id]
f = mapping[0]
e = mapping[1]
order = flex_interp[f]
order_abs = -order
bendingadr = flex_bendingadr[f]
cellnum = flex_cellnum[f]
cx = cellnum[0]
cy = cellnum[1]
cz = cellnum[2]
nstart = flex_nodeadr[f]
edata_base = bendingadr + 1 + e * 10
fe_A = int(flex_bending[edata_base + 0])
fe_B = int(flex_bending[edata_base + 1])
local_A = wp.vec2(flex_bending[edata_base + 2], flex_bending[edata_base + 3])
local_B = wp.vec2(flex_bending[edata_base + 4], flex_bending[edata_base + 5])
stiffness = flex_bending[edata_base + 6]
dn0 = wp.vec3(flex_bending[edata_base + 7], flex_bending[edata_base + 8], flex_bending[edata_base + 9])
if stiffness <= 0.0:
return
# Look up cached face data instead of recomputing
face_id_A = flex_faceadr[f] + fe_A
face_id_B = flex_faceadr[f] + fe_B
quat_A = face_quat_in[worldid, face_id_A]
quat_B = face_quat_in[worldid, face_id_B]
# 1. Compute deformed normals at edge midpoint
t1_A = wp.vec3(0.0)
t2_A = wp.vec3(0.0)
t1_B = wp.vec3(0.0)
t2_B = wp.vec3(0.0)
idx = int(0)
for l0 in range(3):
if l0 > order_abs:
continue
for l1 in range(3):
if l1 > order_abs:
continue
pos_A = face_xpos_in[worldid, face_id_A, idx]
pos_B = face_xpos_in[worldid, face_id_B, idx]
grad0_A = support.flex_dphi(local_A[0], l0, order_abs) * support.flex_phi(local_A[1], l1, order_abs)
grad1_A = support.flex_phi(local_A[0], l0, order_abs) * support.flex_dphi(local_A[1], l1, order_abs)
grad0_B = support.flex_dphi(local_B[0], l0, order_abs) * support.flex_phi(local_B[1], l1, order_abs)
grad1_B = support.flex_phi(local_B[0], l0, order_abs) * support.flex_dphi(local_B[1], l1, order_abs)
t1_A += pos_A * grad0_A
t2_A += pos_A * grad1_A
t1_B += pos_B * grad0_B
t2_B += pos_B * grad1_B
idx += 1
n_A = wp.cross(t1_A, t2_A)
n_B = wp.cross(t1_B, t2_B)
len_A = wp.length(n_A)
len_B = wp.length(n_B)
if len_A < MJ_MINVAL or len_B < MJ_MINVAL:
return
inv_A = 1.0 / len_A
inv_B = 1.0 / len_B
n_A_norm = n_A * inv_A
n_B_norm = n_B * inv_B
# 2. Average face quaternions
if wp.dot(quat_A, quat_B) < 0.0:
quat_B = -quat_B
quat_avg = quat_A + quat_B
quat_avg = wp.normalize(quat_avg)
# rotate dn0 using average quat
dn0_rot = wp.quat_rotate(quat_avg, dn0)
# residual: r = (n_A - n_B) - dn0_rot
r = n_A_norm - n_B_norm - dn0_rot
# 3. Compute projection and cross products
dot_A = wp.dot(n_A_norm, r)
w_A = (r - n_A_norm * dot_A) * inv_A
dot_B = wp.dot(n_B_norm, r)
w_B = (r - n_B_norm * dot_B) * inv_B
wAt2 = wp.cross(w_A, t2_A)
wAt1 = wp.cross(w_A, t1_A)
wBt2 = wp.cross(w_B, t2_B)
wBt1 = wp.cross(w_B, t1_B)
# 4. Apply forces for Face A
_apply_face_forces(
flex_nodebodyid,
flex_face,
xipos_in,
flexnode_xpos_in,
face_id_A,
local_A,
wAt1,
wAt2,
stiffness,
order_abs,
worldid,
flex_spring_body_force_out,
)
# 5. Apply forces for Face B (negative stiffness)
_apply_face_forces(
flex_nodebodyid,
flex_face,
xipos_in,
flexnode_xpos_in,
face_id_B,
local_B,
wBt1,
wBt2,
-stiffness,
order_abs,
worldid,
flex_spring_body_force_out,
)
@event_scope
def passive(m: Model, d: Data):
"""Adds all passive forces."""
@@ -1002,6 +1236,12 @@ def passive(m: Model, d: Data):
],
)
flex_spring_body_force = None
flex_damper_body_force = None
if m.nflex > 0:
flex_spring_body_force = wp.zeros((d.nworld, m.nbody), dtype=wp.spatial_vector, device=d.qfrc_spring.device)
flex_damper_body_force = wp.zeros((d.nworld, m.nbody), dtype=wp.spatial_vector, device=d.qfrc_spring.device)
if not dsbl_spring:
wp.launch(
_flex_elasticity,
@@ -1009,7 +1249,6 @@ def passive(m: Model, d: Data):
inputs=[
m.nflex,
m.opt.timestep,
m.body_dofadr,
m.flex_dim,
m.flex_vertadr,
m.flex_edgeadr,
@@ -1024,19 +1263,22 @@ def passive(m: Model, d: Data):
m.flexedge_length0,
m.flex_stiffness,
m.flex_damping,
d.xipos,
d.flexvert_xpos,
d.flexedge_length,
d.flexedge_velocity,
dsbl_damper,
],
outputs=[d.qfrc_spring],
outputs=[flex_spring_body_force],
)
if not dsbl_spring or not dsbl_damper:
wp.launch(
_flex_bending,
dim=(d.nworld, m.nflexedge),
inputs=[
m.nflex,
m.body_dofadr,
m.body_rootid,
m.flex_dim,
m.flex_vertadr,
m.flex_edgeadr,
@@ -1046,9 +1288,41 @@ def passive(m: Model, d: Data):
m.flex_edge,
m.flex_edgeflap,
m.flex_bending,
m.flex_damping,
d.xipos,
d.subtree_com,
d.flexvert_xpos,
d.cvel,
dsbl_damper,
],
outputs=[d.qfrc_spring],
outputs=[
flex_spring_body_force,
flex_damper_body_force,
],
)
wp.launch(
_flex_passive_bend_interp,
dim=(d.nworld, m.nflexbend_interp),
inputs=[
m.nflex,
m.flex_interp,
m.flex_cellnum,
m.flex_nodeadr,
m.flex_nodenum,
m.flex_bendingadr,
m.flex_nodebodyid,
m.flex_node,
m.flex_bending,
m.flex_centered,
m.flex_faceadr,
m.flex_bend_interp_map,
m.flex_face,
d.xipos,
d.flexnode_xpos,
d.face_xpos,
d.face_quat,
],
outputs=[flex_spring_body_force],
)
gravity_enabled = not (m.opt.disableflags & DisableBit.GRAVITY)
@@ -1082,8 +1356,6 @@ def passive(m: Model, d: Data):
inputs=[
m.nflex,
m.body_rootid,
m.body_dofnum,
m.body_dofadr,
m.flex_interp,
m.flex_cellnum,
m.flex_nodeadr,
@@ -1096,6 +1368,7 @@ def passive(m: Model, d: Data):
m.flex_edgeequality,
m.flex_centered,
m.flex_cell_map,
d.xipos,
d.subtree_com,
d.cvel,
d.flexnode_xpos,
@@ -1103,13 +1376,19 @@ def passive(m: Model, d: Data):
dsbl_damper,
],
outputs=[
d.qfrc_spring,
d.qfrc_damper,
flex_spring_body_force,
flex_damper_body_force,
displ_scratch,
vel_corot_scratch,
],
)
if m.nflex > 0:
if not dsbl_spring:
support.apply_ft(m, d, flex_spring_body_force, d.qfrc_spring, True)
if not dsbl_damper:
support.apply_ft(m, d, flex_damper_body_force, d.qfrc_damper, True)
if m.has_fluid:
_fluid(m, d)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+1 -1
View File
@@ -26,7 +26,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.func
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+16 -8
View File
@@ -39,7 +39,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
# Default value for mat_shininess in MuJoCo is 0.5
# With an 8 bit image format, the maximum value is 255.0
@@ -78,10 +78,16 @@ def sample_texture(
mesh_id: int,
) -> wp.vec3:
uv = wp.vec2(0.0, 0.0)
offset = wp.vec2(0.0, 0.0)
if geom_type[geom_id] == GeomType.PLANE:
local = wp.transpose(rot) @ (hit_point - pos)
uv = wp.vec2(local[0], local[1])
# Replicate MuJoCo's OBJECT_PLANE texgen for planes (render_gl3.c settexture):
# s = 0.5 * texrepeat_x * x - 0.5, t = -0.5 * texrepeat_y * y - 0.5, with (x, y)
# the plane-local hit coordinates. The -0.5 is the texgen w-term, independent of
# texrepeat, so it is applied as an offset after the tex_repeat scale below.
uv = wp.vec2(0.5 * local[0], -0.5 * local[1])
offset = wp.vec2(-0.5, -0.5)
if geom_type[geom_id] == GeomType.MESH:
if f < 0 or mesh_id < 0:
@@ -93,8 +99,8 @@ def sample_texture(
uv2 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][2]]
uv = uv0 * bary_u + uv1 * bary_v + uv2 * (1.0 - bary_u - bary_v)
u = uv[0] * tex_repeat[0]
v = uv[1] * tex_repeat[1]
u = uv[0] * tex_repeat[0] + offset[0]
v = uv[1] * tex_repeat[1] + offset[1]
u = u - wp.floor(u)
v = v - wp.floor(v)
tex_color = wp.texture_sample(tex, wp.vec2(u, v), dtype=wp.vec4)
@@ -545,8 +551,6 @@ def render(m: Model, d: Data, rc: RenderContext):
d: The data on device.
rc: The render context on device.
"""
rc.rgb_data.fill_(rc.background_color)
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
@@ -557,11 +561,11 @@ def render(m: Model, d: Data, rc: RenderContext):
compute_lighting = _make_compute_lighting(cast_ray_first_hit)
# Static parameters extracted for JAX FFI closure.
rc_static = {f.name: getattr(rc, f.name) for f in dataclasses.fields(rc) if f.type in (int, bool, float, wp.vec3)}
rc_static = {f.name: getattr(rc, f.name) for f in dataclasses.fields(rc) if f.type in (int, wp.uint32, bool, float, wp.vec3)}
rc_static["enable_specular_or_emission"] = rc.enable_specular or rc.enable_emission
M_NLIGHT = m.nlight
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"fast_math": rc.use_fast_math})
def _render_megakernel(
# Model:
geom_type: wp.array[int],
@@ -716,6 +720,8 @@ def render(m: Model, d: Data, rc: RenderContext):
# Early Out
if geom_id == -1:
if render_depth[camid]:
depth_out[worldid, depth_adr[camid] + rayid_local] = 0.0
if wp.static(rc_static["render_skybox"]) and render_rgb[camid]:
skybox_id = skybox_tex_id[worldid % skybox_tex_id.shape[0]]
skybox_color = sample_skybox(
@@ -729,6 +735,8 @@ def render(m: Model, d: Data, rc: RenderContext):
skybox_color[2] * 255.0,
255.0,
)
elif render_rgb[camid]:
rgb_out[worldid, rgb_adr[camid] + rayid_local] = wp.static(rc_static["background_color"])
return
if render_depth[camid]:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
Vendored
+1 -1
View File
@@ -19,7 +19,7 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.types import ProjectionType
from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.kernel
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+4 -4
View File
@@ -50,7 +50,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.util_misc import poly_potential
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
@wp.func
@@ -2140,7 +2140,7 @@ def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3:
return wp.spatial_bottom(vec) - wp.cross(dif, wp.spatial_top(vec))
@wp.kernel
@wp.kernel(grid_stride=True)
def _preprocess_tactile_contacts(
# Model:
body_weldid: wp.array[int],
@@ -2468,7 +2468,7 @@ def _contact_match(
@cache_kernel
def _contact_sort(maxmatch: int):
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def contact_sort(
# Model:
sensor_intprm: wp.array2d[int],
@@ -2969,7 +2969,7 @@ def energy_pos(m: Model, d: Data):
@cache_kernel
def _energy_vel_kinetic(nv: int):
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def energy_vel_kinetic(
# Data in:
qvel_in: wp.array2d[float],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
Vendored
+4 -1
View File
@@ -23,7 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import SleepState
from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
# tree_asleep value for fully awake tree
K_AWAKE_VAL = -(1 + types.MJ_MINAWAKE)
@@ -244,6 +244,9 @@ def _wake_tree(
tree_asleep_out: wp.array2d[int],
) -> int:
"""Wakes tree treeid and its associated cycle, returning number of woke trees."""
if treeid < 0 or treeid >= ntree:
return 0
asleep_val = tree_asleep_out[worldid, treeid]
if asleep_val < 0:
if wakeval < asleep_val:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+359 -124
View File
@@ -21,6 +21,8 @@ from mujoco.mjx.third_party.mujoco_warp._src import support
from mujoco.mjx.third_party.mujoco_warp._src import util_misc
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Q_LD_BLOCK_COMPACT
from mujoco.mjx.third_party.mujoco_warp._src.types import Q_LD_BLOCK_SPARSE
from mujoco.mjx.third_party.mujoco_warp._src.types import CamLightType
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
@@ -503,6 +505,101 @@ def kinematics(m: Model, d: Data):
)
@wp.kernel
def _flex_face_kinematics(
# Model:
flex_interp: wp.array[int],
flex_cellnum: wp.array[wp.vec3i],
flex_face_map: wp.array[wp.vec2i],
flex_face: wp.array2d[int],
# Data in:
flexnode_xpos_in: wp.array2d[wp.vec3],
# Data out:
face_xpos_out: wp.array3d[wp.vec3],
face_quat_out: wp.array2d[wp.quat],
):
worldid, face_id = wp.tid()
mapping = flex_face_map[face_id]
f = mapping[0]
face_elem_idx = mapping[1]
order = flex_interp[f]
order_abs = -order
cellnum = flex_cellnum[f]
cx = cellnum[0]
cy = cellnum[1]
cz = cellnum[2]
t1 = wp.vec3(0.0)
t2 = wp.vec3(0.0)
for local_idx in range(9):
gidx = flex_face[face_id, local_idx]
if gidx != -1:
node_pos = flexnode_xpos_in[worldid, gidx]
face_xpos_out[worldid, face_id, local_idx] = node_pos
l0 = local_idx // (order_abs + 1)
l1 = local_idx % (order_abs + 1)
dphi0 = float(l0 - 1) if order_abs == 2 else (-1.0 + 2.0 * float(l0))
dphi1 = float(l1 - 1) if order_abs == 2 else (-1.0 + 2.0 * float(l1))
phi0 = wp.where(l0 == 1, 1.0, 0.0) if order_abs == 2 else 0.5
phi1 = wp.where(l1 == 1, 1.0, 0.0) if order_abs == 2 else 0.5
grad0 = dphi0 * phi1
grad1 = phi0 * dphi1
t1 += node_pos * grad0
t2 += node_pos * grad1
else:
face_xpos_out[worldid, face_id, local_idx] = wp.vec3(0.0)
normal = wp.cross(t1, t2)
normal_axis, _, _, _, _, _ = support.get_face_metadata(cx, cy, cz, face_elem_idx, order_abs)
F = wp.mat33(0.0)
if normal_axis == 0:
F = wp.mat33(
normal[0],
t1[0],
t2[0],
normal[1],
t1[1],
t2[1],
normal[2],
t1[2],
t2[2],
)
elif normal_axis == 1:
F = wp.mat33(
t2[0],
normal[0],
t1[0],
t2[1],
normal[1],
t1[1],
t2[2],
normal[2],
t1[2],
)
else:
F = wp.mat33(
t1[0],
t2[0],
normal[0],
t1[1],
t2[1],
normal[1],
t1[2],
t2[2],
normal[2],
)
face_quat_out[worldid, face_id] = support.mat33_to_quat_polar(F)
@event_scope
def flex(m: Model, d: Data):
# Compute node positions first (needed for interpolated vertex positions)
@@ -569,6 +666,22 @@ def flex(m: Model, d: Data):
],
)
wp.launch(
_flex_face_kinematics,
dim=(d.nworld, m.nflexface),
inputs=[
m.flex_interp,
m.flex_cellnum,
m.flex_face_map,
m.flex_face,
d.flexnode_xpos,
],
outputs=[
d.face_xpos,
d.face_quat,
],
)
@wp.kernel
def _subtree_com_init(
@@ -1108,25 +1221,6 @@ def _qLDiag_div(
D_out[worldid, dofid] = 1.0 / L_in[worldid, diag_i]
@wp.kernel
def _factor_simple(
# Model:
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
# Data in:
M_in: wp.array2d[float],
# In:
simple_dofs: wp.array[int],
# Out:
D_out: wp.array2d[float],
):
# A simple (decoupled) dof's whole factorization is D = 1/M(i,i): no L entries, no elimination.
worldid, s = wp.tid()
dofid = simple_dofs[s]
diag_i = M_rowadr[dofid] + M_rownnz[dofid] - 1
D_out[worldid, dofid] = 1.0 / M_in[worldid, diag_i]
def _factor_i_sparse(m: Model, d: Data, M: wp.array2d[float], L: wp.array2d[float], D: wp.array2d[float]):
"""Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
wp.copy(L, M)
@@ -1138,6 +1232,50 @@ def _factor_i_sparse(m: Model, d: Data, M: wp.array2d[float], L: wp.array2d[floa
wp.launch(_qLDiag_div, dim=(d.nworld, m.nv), inputs=[m.M_rownnz, m.M_rowadr, L], outputs=[D])
@cache_kernel
def _small_cholesky_factorize_block(block_size: int):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
M_rowadr: wp.array[int],
qLD_block_adr: wp.array[int],
# Data in:
M_in: wp.array2d[float],
# In:
block_dof: wp.array[int],
# Out:
D_out: wp.array2d[float],
L_out: wp.array2d[float],
):
worldid, blk = wp.tid()
start = block_dof[blk]
size = wp.static(block_size)
matrix_adr = M_rowadr[start]
factor_adr = qLD_block_adr[start]
if factor_adr == Q_LD_BLOCK_COMPACT:
for i in range(wp.static(block_size)):
D_out[worldid, start + i] = 1.0 / M_in[worldid, matrix_adr + i]
else:
for i in range(wp.static(block_size)):
value = M_in[worldid, matrix_adr + i * (i + 1) // 2 + i]
for k in range(i):
factor = L_out[worldid, factor_adr + k * size + i]
value -= factor * factor
diagonal_value = wp.sqrt(value)
L_out[worldid, factor_adr + i * size + i] = diagonal_value
diagonal_inv = 1.0 / diagonal_value
for j in range(i + 1, size):
value = M_in[worldid, matrix_adr + j * (j + 1) // 2 + i]
for k in range(i):
value -= L_out[worldid, factor_adr + k * size + i] * L_out[worldid, factor_adr + k * size + j]
L_out[worldid, factor_adr + i * size + j] = value * diagonal_inv
return kernel
@cache_kernel
def _tile_cholesky_factorize_block(tile: TileSet):
# One diagonal block of `block_size` dofs per (world, block) tile group. tile_load_indexed gathers
@@ -1173,36 +1311,43 @@ def _tile_cholesky_factorize_block(tile: TileSet):
return kernel
def _factor_block_dense(m: Model, d: Data, M: wp.array2d[float], L: wp.array2d[float]):
def _factor_blocks(
m: Model,
d: Data,
M: wp.array2d[float],
L: wp.array2d[float],
D: wp.array2d[float],
):
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_factorize_block(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, M, tile.elemid, tile.adr],
outputs=[L],
block_dim=m.block_dim.cholesky_factorize,
)
if tile.elemid.size == 0:
wp.launch(
_small_cholesky_factorize_block(tile.size),
dim=(d.nworld, tile.adr.size),
inputs=[m.M_rowadr, m.qLD_block_adr, M, tile.adr],
outputs=[D, L],
block_dim=m.block_dim.small_cholesky,
)
else:
wp.launch_tiled(
_tile_cholesky_factorize_block(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, M, tile.elemid, tile.adr],
outputs=[L],
block_dim=m.block_dim.cholesky_factorize,
)
@event_scope
def factor_m(m: Model, d: Data):
"""Factorization of inertia-like matrix M, assumed spd.
The factor is a per-block decision: dense blocks factor as a packed tile-Cholesky (M_tiles),
sparse blocks via the LDL factor over the LDL region (offset qLD_block_total), and simple
(diagonal) blocks need only D = 1/diag. The passes write disjoint dofs and may all run at once.
Compact blocks use reciprocal diagonals, full small blocks use scalar Cholesky, larger dense
blocks use tile Cholesky, and oversized blocks use sparse LDL.
"""
if m.qLD_has_dense:
_factor_block_dense(m, d, d.M, d.qLD)
if m.qLD_has_sparse:
if m.M_tiles:
_factor_blocks(m, d, d.M, d.qLD, d.qLDiagInv)
if d.qLD.shape[1] > m.qLD_block_total:
_factor_i_sparse(m, d, d.M, d.qLD[:, m.qLD_block_total :], d.qLDiagInv)
if m.qLD_has_simple:
wp.launch(
_factor_simple,
dim=(d.nworld, m.qLD_simple_dofs.size),
inputs=[m.M_rownnz, m.M_rowadr, d.M, m.qLD_simple_dofs],
outputs=[d.qLDiagInv],
)
@wp.kernel
@@ -2158,8 +2303,8 @@ def _transmission(
ten_J_colind: wp.array[int],
actuator_trntype: wp.array[int],
actuator_trnid: wp.array[wp.vec2i],
actuator_gear: wp.array2d[wp.spatial_vector],
actuator_cranklength: wp.array2d[float],
actuator_gear: wp.array2d[wp.spatial_vector],
body_isdofancestor: wp.array2d[int],
# Data in:
qpos_in: wp.array2d[float],
@@ -2771,8 +2916,8 @@ def transmission(m: Model, d: Data):
m.ten_J_colind,
m.actuator_trntype,
m.actuator_trnid,
m.actuator_gear,
m.actuator_cranklength,
m.actuator_gear,
m.body_isdofancestor,
d.qpos,
d.xquat,
@@ -2843,9 +2988,9 @@ def _solve_LD_sparse_fused(nv: int, nlevels: int):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
qLD_block_adr: wp.array[int],
# In:
dof_dense: wp.array[int],
dof_simple: wp.array[int],
L: wp.array2d[float],
D: wp.array2d[float],
all_updates: wp.array[wp.vec3i],
@@ -2859,10 +3004,9 @@ def _solve_LD_sparse_fused(nv: int, nlevels: int):
NLEVELS = wp.static(nlevels)
BLOCK_DIM = wp.block_dim()
# Copy y to x_out for sparse-block dofs only; dense blocks use the packed pass and simple
# (diagonal) blocks use the dedicated 1/diag solve.
# Copy y to x_out for sparse-block dofs only.
for dofid in range(tid, NV, BLOCK_DIM):
if dof_dense[dofid] == 0 and dof_simple[dofid] == 0:
if qLD_block_adr[dofid] == Q_LD_BLOCK_SPARSE:
x_out[worldid, dofid] = y[worldid, dofid]
_syncthreads()
@@ -2880,7 +3024,7 @@ def _solve_LD_sparse_fused(nv: int, nlevels: int):
# Diagonal multiply (sparse-block dofs only)
for dofid in range(tid, NV, BLOCK_DIM):
if dof_dense[dofid] == 0 and dof_simple[dofid] == 0:
if qLD_block_adr[dofid] == Q_LD_BLOCK_SPARSE:
x_out[worldid, dofid] *= D[worldid, dofid]
_syncthreads()
@@ -2918,48 +3062,63 @@ def _solve_LD_sparse(
wp.launch(
_solve_LD_sparse_fused(m.nv, nlevels),
dim=(d.nworld, dim_block),
inputs=[m.qLD_dof_dense, m.qLD_dof_simple, L, D, m.qLD_all_updates, m.qLD_level_offsets, y],
inputs=[m.qLD_block_adr, L, D, m.qLD_all_updates, m.qLD_level_offsets, y],
outputs=[x],
block_dim=dim_block,
)
@wp.kernel
def _solve_simple(
@wp.func
def _small_cholesky_solve(
# In:
simple_dofs: wp.array[int],
D: wp.array2d[float],
y: wp.array2d[float],
block_size: int,
worldid: int,
factor_adr: int,
start: int,
L_in: wp.array2d[float],
y_in: wp.array2d[float],
# Out:
x_out: wp.array2d[float],
):
# A simple (decoupled) dof's solve is just x = (1/diag) * y.
worldid, s = wp.tid()
dofid = simple_dofs[s]
x_out[worldid, dofid] = D[worldid, dofid] * y[worldid, dofid]
for i in range(block_size):
value = y_in[worldid, start + i]
for k in range(i):
value -= L_in[worldid, factor_adr + k * block_size + i] * x_out[worldid, start + k]
x_out[worldid, start + i] = value / L_in[worldid, factor_adr + i * block_size + i]
for reverse_i in range(block_size):
i = block_size - 1 - reverse_i
value = x_out[worldid, start + i]
for k in range(i + 1, block_size):
value -= L_in[worldid, factor_adr + i * block_size + k] * x_out[worldid, start + k]
x_out[worldid, start + i] = value / L_in[worldid, factor_adr + i * block_size + i]
@wp.kernel
def _factor_solve_simple(
# Model:
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
# Data in:
M_in: wp.array2d[float],
# In:
simple_dofs: wp.array[int],
y: wp.array2d[float],
# Out:
D_out: wp.array2d[float],
x_out: wp.array2d[float],
):
# Fused factor+solve for a simple dof: read M(i,i) once, emit D = 1/diag and x = D * y.
worldid, s = wp.tid()
dofid = simple_dofs[s]
diag_i = M_rowadr[dofid] + M_rownnz[dofid] - 1
d_inv = 1.0 / M_in[worldid, diag_i]
D_out[worldid, dofid] = d_inv
x_out[worldid, dofid] = d_inv * y[worldid, dofid]
@cache_kernel
def _small_cholesky_solve_block(block_size: int):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
qLD_block_adr: wp.array[int],
# In:
block_dof: wp.array[int],
D_in: wp.array2d[float],
L_in: wp.array2d[float],
y_in: wp.array2d[float],
# Out:
x_out: wp.array2d[float],
):
worldid, blk = wp.tid()
start = block_dof[blk]
size = wp.static(block_size)
factor_adr = qLD_block_adr[start]
if factor_adr == Q_LD_BLOCK_COMPACT:
for i in range(wp.static(block_size)):
x_out[worldid, start + i] = D_in[worldid, start + i] * y_in[worldid, start + i]
else:
_small_cholesky_solve(size, worldid, factor_adr, start, L_in, y_in, x_out)
return kernel
@cache_kernel
@@ -2992,18 +3151,34 @@ def _tile_cholesky_solve_block(tile: TileSet):
return kernel
def _solve_block_dense(m: Model, d: Data, L: wp.array2d[float], x: wp.array2d[float], y: wp.array2d[float]):
def _solve_blocks(
m: Model,
d: Data,
L: wp.array2d[float],
D: wp.array2d[float],
x: wp.array2d[float],
y: wp.array2d[float],
):
for tile in m.M_tiles:
# The triangular back-substitution is largely sequential, so large blocks prefer fewer threads
# for better occupancy while moderate blocks still want a couple warps (16/27->64, 60->32).
block_dim = m.block_dim.cholesky_solve if tile.size <= 40 else 32
wp.launch_tiled(
_tile_cholesky_solve_block(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, tile.adr, L, y],
outputs=[x],
block_dim=block_dim,
)
if tile.elemid.size == 0:
wp.launch(
_small_cholesky_solve_block(tile.size),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, tile.adr, D, L, y],
outputs=[x],
block_dim=m.block_dim.small_cholesky,
)
else:
# The triangular back-substitution is largely sequential, so large blocks prefer fewer threads
# for better occupancy while moderate blocks still want a couple warps (16/27->64, 60->32).
block_dim = m.block_dim.cholesky_solve if tile.size <= 40 else 32
wp.launch_tiled(
_tile_cholesky_solve_block(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, tile.adr, L, y],
outputs=[x],
block_dim=block_dim,
)
def solve_LD(
@@ -3016,25 +3191,21 @@ def solve_LD(
):
"""Computes backsubstitution for the inertia factorization.
The choice is per-block. Dense blocks back-substitute from the packed Cholesky region of L; sparse
blocks from the LDL region (offset qLD_block_total); simple (diagonal) blocks are a plain x = D*y.
The passes write disjoint dofs; the sparse pass skips dense and simple dofs so it does not clobber
their results.
Compact blocks use reciprocal diagonals, full small blocks use scalar Cholesky, dense blocks use
tile Cholesky, and sparse blocks use the LDL region.
Args:
m: The model containing factorization and sparsity information.
d: The data object containing workspace and factorization results.
L: The factor: packed dense region followed by the nC LDL region.
D: Diagonal factor (1/diag) for the sparse LDL and simple regions.
D: Reciprocal diagonal for compact and sparse blocks.
x: Output array for the solution.
y: Input right-hand side array.
"""
if m.qLD_has_dense:
_solve_block_dense(m, d, L, x, y)
if m.qLD_has_sparse:
if m.M_tiles:
_solve_blocks(m, d, L, D, x, y)
if L.shape[1] > m.qLD_block_total:
_solve_LD_sparse(m, d, L[:, m.qLD_block_total :], D, x, y)
if m.qLD_has_simple:
wp.launch(_solve_simple, dim=(d.nworld, m.qLD_simple_dofs.size), inputs=[m.qLD_simple_dofs, D, y], outputs=[x])
@event_scope
@@ -3089,49 +3260,113 @@ def _tile_cholesky_factorize_solve_block(tile: TileSet):
return kernel
def _factor_solve_block_dense(
m: Model, d: Data, M: wp.array2d[float], x: wp.array2d[float], y: wp.array2d[float], L: wp.array2d[float]
@cache_kernel
def _small_cholesky_factorize_solve_block(block_size: int):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
M_rowadr: wp.array[int],
qLD_block_adr: wp.array[int],
# Data in:
M_in: wp.array2d[float],
# In:
block_dof: wp.array[int],
y: wp.array2d[float],
# Out:
D_out: wp.array2d[float],
x_out: wp.array2d[float],
L_out: wp.array2d[float],
):
worldid, blk = wp.tid()
start = block_dof[blk]
size = wp.static(block_size)
matrix_adr = M_rowadr[start]
factor_adr = qLD_block_adr[start]
if factor_adr == Q_LD_BLOCK_COMPACT:
for i in range(wp.static(block_size)):
inverse = 1.0 / M_in[worldid, matrix_adr + i]
D_out[worldid, start + i] = inverse
x_out[worldid, start + i] = inverse * y[worldid, start + i]
else:
for i in range(wp.static(block_size)):
diagonal_value = M_in[worldid, matrix_adr + i * (i + 1) // 2 + i]
rhs_value = y[worldid, start + i]
for k in range(i):
factor = L_out[worldid, factor_adr + k * size + i]
diagonal_value -= factor * factor
rhs_value -= factor * x_out[worldid, start + k]
diagonal_factor = wp.sqrt(diagonal_value)
L_out[worldid, factor_adr + i * size + i] = diagonal_factor
diagonal_inv = 1.0 / diagonal_factor
x_out[worldid, start + i] = rhs_value * diagonal_inv
for j in range(i + 1, size):
value = M_in[worldid, matrix_adr + j * (j + 1) // 2 + i]
for k in range(i):
value -= L_out[worldid, factor_adr + k * size + i] * L_out[worldid, factor_adr + k * size + j]
L_out[worldid, factor_adr + i * size + j] = value * diagonal_inv
for reverse_i in range(wp.static(block_size)):
i = size - 1 - reverse_i
value = x_out[worldid, start + i]
for k in range(i + 1, size):
value -= L_out[worldid, factor_adr + i * size + k] * x_out[worldid, start + k]
x_out[worldid, start + i] = value / L_out[worldid, factor_adr + i * size + i]
return kernel
def _factor_solve_blocks(
m: Model,
d: Data,
M: wp.array2d[float],
L: wp.array2d[float],
D: wp.array2d[float],
x: wp.array2d[float],
y: wp.array2d[float],
):
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_factorize_solve_block(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, M, tile.elemid, tile.adr, y],
outputs=[x, L],
block_dim=m.block_dim.cholesky_factorize_solve,
)
if tile.elemid.size == 0:
wp.launch(
_small_cholesky_factorize_solve_block(tile.size),
dim=(d.nworld, tile.adr.size),
inputs=[m.M_rowadr, m.qLD_block_adr, M, tile.adr, y],
outputs=[D, x, L],
block_dim=m.block_dim.small_cholesky,
)
else:
wp.launch_tiled(
_tile_cholesky_factorize_solve_block(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.qLD_block_adr, M, tile.elemid, tile.adr, y],
outputs=[x, L],
block_dim=m.block_dim.cholesky_factorize_solve,
)
def factor_solve_i(m, d, M, L, D, x, y):
"""Factorizes and solves the inertia-like linear system.
The choice is per-block (see factor_m): dense blocks factor+solve via the packed Cholesky, sparse
blocks via the LDL region, simple (diagonal) blocks via D = 1/diag. Factorizes M, solves for x.
Compact blocks use reciprocal diagonals, full small blocks use scalar Cholesky, dense blocks use
tile Cholesky, and sparse blocks use LDL. Factorizes M and solves for x.
Args:
m: The model containing factorization and sparsity information.
d: The data object containing workspace and factorization results.
M: The inertia-like matrix to factorize (CSR, length nC).
L: Output factor: packed dense region followed by the nC LDL region (sized like d.qLD).
D: Output diagonal factor (1/diag) for the sparse LDL and simple regions.
D: Output reciprocal diagonal for compact and sparse blocks.
x: Output array for the solution.
y: Input right-hand side array.
"""
# Per-block: dense blocks factor+solve via the packed Cholesky; sparse blocks via the LDL region
# (offset qLD_block_total); simple blocks via 1/diag. The passes write disjoint dofs.
if m.qLD_has_dense:
_factor_solve_block_dense(m, d, M, x, y, L)
if m.qLD_has_sparse:
if m.M_tiles:
_factor_solve_blocks(m, d, M, L, D, x, y)
if L.shape[1] > m.qLD_block_total:
L_ldl = L[:, m.qLD_block_total :]
_factor_i_sparse(m, d, M, L_ldl, D)
_solve_LD_sparse(m, d, L_ldl, D, x, y)
if m.qLD_has_simple:
wp.launch(
_factor_solve_simple,
dim=(d.nworld, m.qLD_simple_dofs.size),
inputs=[m.M_rownnz, m.M_rowadr, M, m.qLD_simple_dofs, y],
outputs=[D, x],
)
@cache_kernel
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+1243 -1058
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+233 -38
View File
@@ -26,11 +26,11 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import State
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
# TODO(team): kernel analyzer array slice?
@@ -39,7 +39,7 @@ def next_act(
# Model:
opt_timestep: float, # kernel_analyzer: ignore
actuator_dyntype: int, # kernel_analyzer: ignore
actuator_dynprm: vec10f, # kernel_analyzer: ignore
actuator_dynprm: vec10, # kernel_analyzer: ignore
actuator_actrange: wp.vec2, # kernel_analyzer: ignore
# Data In:
act_in: float, # kernel_analyzer: ignore
@@ -67,7 +67,7 @@ def next_act(
@wp.func
def mat33_to_quat_polar(F: wp.mat33) -> wp.quat:
cell_quat = wp.quat(0.0, 0.0, 0.0, 1.0)
for _iter in range(10):
for _iter in range(50):
rot = wp.quat_to_matrix(cell_quat)
rot_t = wp.transpose(rot)
col1_rot = rot_t[0]
@@ -129,35 +129,12 @@ def compute_interp_cell_quat(
node_pos = flexnode_xpos_in[worldid, nstart + gidx]
if order == 1:
dphi_x = float(-1) if li == 0 else float(1)
dphi_y = float(-1) if lj == 0 else float(1)
dphi_z = float(-1) if lk == 0 else float(1)
phi_x = float(0.5)
phi_y = float(0.5)
phi_z = float(0.5)
else:
if li == 0:
dphi_x = -1.0
elif li == 1:
dphi_x = 0.0
else:
dphi_x = 1.0
if lj == 0:
dphi_y = -1.0
elif lj == 1:
dphi_y = 0.0
else:
dphi_y = 1.0
if lk == 0:
dphi_z = -1.0
elif lk == 1:
dphi_z = 0.0
else:
dphi_z = 1.0
phi_x = 0.5 if li == 0 or li == 2 else 1.0
phi_y = 0.5 if lj == 0 or lj == 2 else 1.0
phi_z = 0.5 if lk == 0 or lk == 2 else 1.0
dphi_x = float(-1) if li == 0 else float(1)
dphi_y = float(-1) if lj == 0 else float(1)
dphi_z = float(-1) if lk == 0 else float(1)
phi_x = float(0.5)
phi_y = float(0.5)
phi_z = float(0.5)
grad_x = dphi_x * phi_y * phi_z
grad_y = phi_x * dphi_y * phi_z
@@ -175,7 +152,7 @@ def compute_interp_cell_quat(
@cache_kernel
def mul_m_kernel(check_skip: bool):
@wp.kernel(module="unique")
@wp.kernel(module="unique", grid_stride=False)
def _mul_m(
# Model:
M_mulm_rowadr: wp.array[int],
@@ -212,7 +189,7 @@ def mul_m_kernel(check_skip: bool):
@cache_kernel
def mul_m_dense(nv: int, check_skip: bool):
@wp.kernel(module="unique")
@wp.kernel(module="unique", grid_stride=False)
def _mul_m_dense(
# Data in:
M_in: wp.array3d[float], # kernel_analyzer: ignore
@@ -550,7 +527,7 @@ def jac_dof(
@cache_kernel
def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def _jac(
# Model:
body_parentid: wp.array[int],
@@ -702,7 +679,7 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
if sig >= (1 << State.NSTATE):
raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def _get_state(
# Model:
nq: int,
@@ -857,7 +834,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
if sig >= (1 << State.NSTATE):
raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
@wp.kernel(module="unique", enable_backward=False)
@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def _set_state(
# Model:
nq: int,
@@ -1083,3 +1060,221 @@ def select_top4_weights(
)
return selected_b, selected_W
@wp.func
def get_face_metadata(
# In:
cellnum_x: int,
cellnum_y: int,
cellnum_z: int,
face_elem_idx: int,
order_abs: int,
) -> Tuple[int, int, int, int, int, int]:
size01 = cellnum_y * cellnum_z
size23 = cellnum_x * cellnum_z
size45 = cellnum_x * cellnum_y
face_id = 0
within_face = 0
if face_elem_idx < size01:
face_id = 0
within_face = face_elem_idx
elif face_elem_idx < 2 * size01:
face_id = 1
within_face = face_elem_idx - size01
elif face_elem_idx < 2 * size01 + size23:
face_id = 2
within_face = face_elem_idx - 2 * size01
elif face_elem_idx < 2 * size01 + 2 * size23:
face_id = 3
within_face = face_elem_idx - 2 * size01 - size23
elif face_elem_idx < 2 * size01 + 2 * size23 + size45:
face_id = 4
within_face = face_elem_idx - 2 * size01 - 2 * size23
else:
face_id = 5
within_face = face_elem_idx - 2 * size01 - 2 * size23 - size45
normal_axis = face_id // 2
c1 = 0
if face_id == 0 or face_id == 1:
c1 = cellnum_z
elif face_id == 2 or face_id == 3:
c1 = cellnum_x
else:
c1 = cellnum_y
fixed_dim = 0
if normal_axis == 0:
fixed_dim = cellnum_x
elif normal_axis == 1:
fixed_dim = cellnum_y
else:
fixed_dim = cellnum_z
g_fixed = (face_id % 2) * fixed_dim * order_abs
q0 = within_face // c1
q1 = within_face % c1
ny_g = cellnum_y * order_abs + 1
nz_g = cellnum_z * order_abs + 1
return normal_axis, g_fixed, q0, q1, ny_g, nz_g
@wp.func
def gather_face_node_index_fast(
# In:
normal_axis: int,
g_fixed: int,
q0: int,
q1: int,
ny_g: int,
nz_g: int,
local_idx: int,
order_abs: int,
) -> int:
l0 = local_idx // (order_abs + 1)
l1 = local_idx % (order_abs + 1)
g = wp.vec3i(0, 0, 0)
if normal_axis == 0:
g = wp.vec3i(g_fixed, q0 * order_abs + l0, q1 * order_abs + l1)
elif normal_axis == 1:
g = wp.vec3i(q1 * order_abs + l1, g_fixed, q0 * order_abs + l0)
else:
g = wp.vec3i(q0 * order_abs + l0, q1 * order_abs + l1, g_fixed)
return g[0] * ny_g * nz_g + g[1] * nz_g + g[2]
@wp.func
def gather_face_node_index(
# In:
cellnum_x: int,
cellnum_y: int,
cellnum_z: int,
face_elem_idx: int,
local_idx: int,
order_abs: int,
) -> int:
normal_axis, g_fixed, q0, q1, ny_g, nz_g = get_face_metadata(cellnum_x, cellnum_y, cellnum_z, face_elem_idx, order_abs)
return gather_face_node_index_fast(normal_axis, g_fixed, q0, q1, ny_g, nz_g, local_idx, order_abs)
@wp.func
def compute_interp_face_quat(
# Data in:
flexnode_xpos_in: wp.array2d[wp.vec3],
# In:
cellnum_x: int,
cellnum_y: int,
cellnum_z: int,
face_elem_idx: int,
nstart: int,
order_abs: int,
worldid: int,
) -> wp.quat:
normal_axis, g_fixed, q0, q1, ny_g, nz_g = get_face_metadata(
cellnum_x,
cellnum_y,
cellnum_z,
face_elem_idx,
order_abs,
)
t1 = wp.vec3(0.0)
t2 = wp.vec3(0.0)
npc = (order_abs + 1) * (order_abs + 1)
for local_idx in range(9):
if local_idx < npc:
gidx = gather_face_node_index_fast(
normal_axis,
g_fixed,
q0,
q1,
ny_g,
nz_g,
local_idx,
order_abs,
)
node_pos = flexnode_xpos_in[worldid, nstart + gidx]
l0 = local_idx // (order_abs + 1)
l1 = local_idx % (order_abs + 1)
dphi0 = -1.0 + 2.0 * float(l0)
dphi1 = -1.0 + 2.0 * float(l1)
phi0 = 0.5
phi1 = 0.5
grad0 = dphi0 * phi1
grad1 = phi0 * dphi1
t1 += node_pos * grad0
t2 += node_pos * grad1
normal = wp.cross(t1, t2)
F = wp.mat33(0.0)
if normal_axis == 0:
F = wp.mat33(
normal[0],
t1[0],
t2[0],
normal[1],
t1[1],
t2[1],
normal[2],
t1[2],
t2[2],
)
elif normal_axis == 1:
F = wp.mat33(
t2[0],
normal[0],
t1[0],
t2[1],
normal[1],
t1[1],
t2[2],
normal[2],
t1[2],
)
else:
F = wp.mat33(
t1[0],
t2[0],
normal[0],
t1[1],
t2[1],
normal[1],
t1[2],
t2[2],
normal[2],
)
return mat33_to_quat_polar(F)
@wp.func
def flex_phi(s: float, i: int, order: int) -> float:
return 1.0 - s if i == 0 else s
@wp.func
def flex_dphi(s: float, i: int, order: int) -> float:
return -1.0 if i == 0 else 1.0
@wp.func
def dphi2D(s0: float, l0: int, s1: float, l1: int, order: int, direction: int) -> float:
if direction == 0:
return flex_dphi(s0, l0, order) * flex_phi(s1, l1, order)
else:
return flex_phi(s0, l0, order) * flex_dphi(s1, l1, order)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py
Vendored
-351
View File
@@ -1,351 +0,0 @@
# Copyright 2025 The Newton Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for flex."""
import mujoco
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
import mujoco.mjx.third_party.mujoco_warp as mjw
from mujoco.mjx.third_party.mujoco_warp import test_data
_TRILINEAR_STRAIN_XML = """
<mujoco>
<option gravity="0 0 -9.81">
<flag contact="disable"/>
</option>
<worldbody>
<flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
dof="trilinear">
<edge equality="strain"/>
<contact selfcollide="none"/>
</flexcomp>
</worldbody>
</mujoco>
"""
# tolerance for difference between MuJoCo and MJWarp, mostly due to float32
_TOLERANCE = 5e-4
class TrilinearFlexTest(parameterized.TestCase):
def test_flexstrain_constraint_at_rest(self):
"""Test FLEXSTRAIN constraint count and residuals match MuJoCo at rest."""
mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
mjd = mujoco.MjData(mjm)
mujoco.mj_forward(mjm, mjd)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)
mjw.fwd_position(m, d)
mjw.make_constraint(m, d)
# constraint counts should match
ne_warp = d.ne.numpy()[0]
self.assertEqual(ne_warp, mjd.ne, f"ne mismatch: warp={ne_warp}, mj={mjd.ne}")
nefc_warp = d.nefc.numpy()[0]
self.assertEqual(nefc_warp, mjd.nefc, f"nefc mismatch: warp={nefc_warp}, mj={mjd.nefc}")
# residuals should match
efc_pos = d.efc.pos.numpy()[0, :nefc_warp]
efc_pos_mj = mjd.efc_pos[: mjd.nefc]
np.testing.assert_allclose(efc_pos, efc_pos_mj, atol=1e-5, err_msg="FLEXSTRAIN residuals should match MuJoCo at rest")
def test_flexstrain_constraint_perturbed(self):
"""Test FLEXSTRAIN residuals and Jacobians match MuJoCo under perturbation."""
mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
mjd = mujoco.MjData(mjm)
# perturb first node
mjd.qpos[0] += 0.01
mjd.qpos[3] += 0.005
mujoco.mj_forward(mjm, mjd)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)
mjw.fwd_position(m, d)
mjw.make_constraint(m, d)
nefc = d.nefc.numpy()[0]
nv = mjm.nv
self.assertEqual(nefc, mjd.nefc)
# residuals
efc_pos_warp = d.efc.pos.numpy()[0, :nefc]
efc_pos_mj = mjd.efc_pos[: mjd.nefc]
np.testing.assert_allclose(
efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under perturbation"
)
# Jacobians
if mujoco.mj_isSparse(mjm):
mj_efc_J = np.zeros((mjd.nefc, nv))
mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind)
else:
mj_efc_J = mjd.efc_J.reshape((mjd.nefc, nv))
if m.is_sparse:
warp_efc_J = np.zeros((nefc, nv))
mujoco.mju_sparse2dense(
warp_efc_J,
d.efc.J.numpy()[0, 0],
d.efc.J_rownnz.numpy()[0, :nefc],
d.efc.J_rowadr.numpy()[0, :nefc],
d.efc.J_colind.numpy()[0, 0],
)
else:
warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv]
np.testing.assert_allclose(warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo")
def test_flexstrain_constraint_rotated(self):
"""Test FLEXSTRAIN residuals and Jacobians match MuJoCo under large rotation perturbation."""
mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
mjd = mujoco.MjData(mjm)
# Apply a rotation perturbation: rotate all node positions around Y axis by 30 degrees
# (0.5235 radians)
theta = 0.5235
cos_t = np.cos(theta)
sin_t = np.sin(theta)
for i in range(0, mjm.nq, 3):
x = mjd.qpos[i]
z = mjd.qpos[i + 2]
mjd.qpos[i] = x * cos_t - z * sin_t
mjd.qpos[i + 2] = x * sin_t + z * cos_t
mujoco.mj_forward(mjm, mjd)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)
mjw.fwd_position(m, d)
mjw.make_constraint(m, d)
nefc = d.nefc.numpy()[0]
nv = mjm.nv
self.assertEqual(nefc, mjd.nefc)
# residuals
efc_pos_warp = d.efc.pos.numpy()[0, :nefc]
efc_pos_mj = mjd.efc_pos[: mjd.nefc]
np.testing.assert_allclose(
efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under rotation"
)
# Jacobians
if mujoco.mj_isSparse(mjm):
mj_efc_J = np.zeros((mjd.nefc, nv))
mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind)
else:
mj_efc_J = mjd.efc_J.reshape((mjd.nefc, nv))
if m.is_sparse:
warp_efc_J = np.zeros((nefc, nv))
mujoco.mju_sparse2dense(
warp_efc_J,
d.efc.J.numpy()[0, 0],
d.efc.J_rownnz.numpy()[0, :nefc],
d.efc.J_rowadr.numpy()[0, :nefc],
d.efc.J_colind.numpy()[0, 0],
)
else:
warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv]
np.testing.assert_allclose(
warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo under rotation"
)
def test_flexstrain_rotational_invariance(self):
"""Test that FLEXSTRAIN residuals are invariant under rigid translation."""
mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
mjd = mujoco.MjData(mjm)
mujoco.mj_forward(mjm, mjd)
# Get reference residuals
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)
mjw.fwd_position(m, d)
mjw.make_constraint(m, d)
nefc = d.nefc.numpy()[0]
efc_pos_rest = d.efc.pos.numpy()[0, :nefc].copy()
# Apply uniform translation to all nodes (rigid motion)
mjd2 = mujoco.MjData(mjm)
# All flex nodes have 3 DOFs (slide joints), shift all x by 0.1
for i in range(0, mjm.nq, 3):
mjd2.qpos[i] += 0.1 # shift x
mujoco.mj_forward(mjm, mjd2)
d2 = mjw.put_data(mjm, mjd2)
mjw.fwd_position(m, d2)
mjw.make_constraint(m, d2)
nefc2 = d2.nefc.numpy()[0]
efc_pos_shifted = d2.efc.pos.numpy()[0, :nefc2]
# Residuals should remain near zero for rigid translation
np.testing.assert_allclose(
efc_pos_shifted, efc_pos_rest, atol=1e-4, err_msg="FLEXSTRAIN residuals should be invariant under rigid translation"
)
def test_trilinear_gravity_parity(self):
"""Test that trilinear simulation matches MuJoCo after multiple steps."""
mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
# MuJoCo reference
mjd = mujoco.MjData(mjm)
for _ in range(10):
mujoco.mj_step(mjm, mjd)
# Warp
mjd_warp = mujoco.MjData(mjm)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd_warp)
z0 = d.qpos.numpy()[0, 2]
for _ in range(10):
mjw.step(m, d)
qpos_warp = d.qpos.numpy()[0]
qpos_mj = mjd.qpos
# The cube should have fallen
self.assertLess(qpos_warp[2], z0, "Cube should fall under gravity")
# Allow larger tolerance for accumulated integration error
np.testing.assert_allclose(qpos_warp, qpos_mj, atol=0.01, err_msg="Trilinear qpos diverges from MuJoCo after 10 steps")
def test_trilinear_node_positions(self):
"""Test that flexnode_xpos are computed correctly from body kinematics."""
mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
mjd = mujoco.MjData(mjm)
mujoco.mj_forward(mjm, mjd)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)
mjw.fwd_position(m, d)
# Compute expected node positions: xpos_n = body_xpos + body_xmat @ flex_node
nflexnode = mjm.nflexnode
warp_xpos = d.flexnode_xpos.numpy()[0, :nflexnode]
nodeadr = mjm.flex_nodeadr[0]
nodenum = mjm.flex_nodenum[0]
for n in range(nodenum):
bodyid = mjm.flex_nodebodyid[nodeadr + n]
body_xpos = mjd.xpos[bodyid]
body_xmat = mjd.xmat[bodyid].reshape(3, 3)
node_local = mjm.flex_node[nodeadr + n]
expected = body_xpos + body_xmat @ node_local
np.testing.assert_allclose(warp_xpos[n], expected, atol=1e-5, err_msg=f"flexnode_xpos mismatch for node {n}")
def test_trilinear_passive_forces_parity(self):
"""Test passive forces (elasticity) match MuJoCo for trilinear flex."""
xml = """
<mujoco>
<worldbody>
<flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
dof="trilinear">
<elasticity young="1e4" poisson="0.1" damping="0.01"/>
<contact selfcollide="none"/>
</flexcomp>
</worldbody>
</mujoco>
"""
mjm = mujoco.MjModel.from_xml_string(xml)
mjd = mujoco.MjData(mjm)
# perturb first node to generate non-zero elasticity forces
mjd.qpos[0] += 0.01
mjd.qpos[3] += 0.005
mujoco.mj_forward(mjm, mjd)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)
mjw.fwd_position(m, d)
mjw.passive(m, d)
qfrc_passive_warp = d.qfrc_passive.numpy()[0]
qfrc_passive_mj = mjd.qfrc_passive
# Verify they match
np.testing.assert_allclose(
qfrc_passive_warp, qfrc_passive_mj, atol=_TOLERANCE, err_msg="qfrc_passive mismatch for trilinear flex with elasticity"
)
@parameterized.parameters("strain", "true")
def test_trilinear_equality_types(self, equality):
"""Test trilinear with different equality types."""
xml = f"""
<mujoco>
<option gravity="0 0 -9.81">
<flag contact="disable"/>
</option>
<worldbody>
<flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
dof="trilinear">
<edge equality="{equality}"/>
<contact selfcollide="none"/>
</flexcomp>
</worldbody>
</mujoco>
"""
mjm, mjd, m, d = test_data.fixture(xml=xml)
# Should not crash
mjw.forward(m, d)
# Constraint count should match
self.assertEqual(d.nefc.numpy()[0], mjd.nefc)
def test_trilinear_contact_qfrc_constraint(self):
"""Test qfrc_constraint parity for trilinear flex with ground contacts."""
xml = """
<mujoco>
<option gravity="0 0 -9.81"/>
<worldbody>
<geom type="plane" size="1 1 0.1"/>
<flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
pos="0 0 0.05" name="cube" dim="3" mass="1" radius="0.02"
dof="trilinear">
<edge equality="strain"/>
<contact selfcollide="none"/>
</flexcomp>
</worldbody>
</mujoco>
"""
mjm, mjd, m, d = test_data.fixture(xml=xml)
mjw.forward(m, d)
# Verify contacts are generated
nacon = d.nacon.numpy()[0]
self.assertGreater(nacon, 0, "Expected contacts between flex and plane")
self.assertEqual(nacon, mjd.ncon)
# Verify qfrc_constraint parity
qfrc_warp = d.qfrc_constraint.numpy()[0]
qfrc_mj = mjd.qfrc_constraint
np.testing.assert_allclose(qfrc_warp, qfrc_mj, atol=1e-4, err_msg="qfrc_constraint mismatch for trilinear flex contacts")
if __name__ == "__main__":
absltest.main()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+156 -118
View File
@@ -14,7 +14,7 @@
# ==============================================================================
import dataclasses
import enum
from typing import Callable
from typing import Callable, Optional
import mujoco
import numpy as np
@@ -38,6 +38,13 @@ TILE_SIZE_JTDAJ_DENSE = 16
# max M block size where dense tile-Cholesky beats sparse LDL (wins to ~64, degrades past ~80)
M_BLOCK_DENSE_MAX = 64
# max M block size where scalar Cholesky beats dense tile-Cholesky
M_BLOCK_SCALAR_MAX = 6
# qLD_block_adr sentinels for factorless block layouts
Q_LD_BLOCK_COMPACT = -2
Q_LD_BLOCK_SPARSE = -1
# maximum number of plugin attributes
_NPLUGINATTR = 128
@@ -59,6 +66,7 @@ class BlockDim:
cholesky_factorize: block-dense Cholesky factorize block dimension (smooth)
cholesky_factorize_solve: block-dense Cholesky factorize+solve block dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
small_cholesky: scalar small-block Cholesky block dimension (smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
@@ -89,6 +97,7 @@ class BlockDim:
cholesky_factorize: int = 32
cholesky_factorize_solve: int = 32
cholesky_solve: int = 64
small_cholesky: int = 64
solve_LD_sparse_fused: int = 128
# solver
update_gradient_cholesky: int = 64
@@ -149,6 +158,7 @@ class OverflowType(enum.IntFlag):
HFIELD: height field collision overflow
CONTACT_MATCH: contact match sensor overflow
NVMAX: nvmax overflow (islands)
EPA_HORIZON: EPA horizon buffer overflow
"""
NEFC = 1 << 0
@@ -159,6 +169,7 @@ class OverflowType(enum.IntFlag):
HFIELD = 1 << 5
CONTACT_MATCH = 1 << 6
NVMAX = 1 << 7
EPA_HORIZON = 1 << 8
class CamLightType(enum.IntEnum):
@@ -760,6 +771,14 @@ class vec8i(wp.types.vector(length=8, dtype=int)):
pass
class vec16f(wp.types.vector(length=16, dtype=float)):
pass
class vec16i(wp.types.vector(length=16, dtype=int)):
pass
class vec10f(wp.types.vector(length=10, dtype=float)):
pass
@@ -789,6 +808,7 @@ vec6 = vec6f
vec8 = vec8f
vec10 = vec10f
vec11 = vec11f
vec16 = vec16f
vec128 = vec_pluginattr
mat23 = mat23f
mat43 = mat43f
@@ -915,12 +935,12 @@ class TileSet:
Attributes:
adr: address of each tile in the set
size: size of all the tiles in this set
elemid: flat per-block gather indices into CSR M for tile_load_indexed (absent -> nC sentinel)
elemid: flat CSR gather indices for tiled blocks; empty selects the native-layout scalar path
"""
adr: wp.array[int]
size: int
elemid: wp.array[int] = None
elemid: wp.array[int] = dataclasses.field(default_factory=lambda: wp.array([], dtype=int))
def __eq__(self, other) -> bool:
if self.__class__ is not other.__class__:
@@ -1026,6 +1046,7 @@ class Model:
body_treeid: id of body's tree; -1: static (nbody,)
body_geomnum: number of geoms (nbody,)
body_geomadr: start addr of geoms; -1: no geoms (nbody,)
body_simple: body simple type (nbody,)
body_pos: position offset rel. to parent body (*, nbody, 3)
body_quat: orientation offset rel. to parent body (*, nbody, 4)
body_ipos: local position of center of mass (*, nbody, 3)
@@ -1037,6 +1058,7 @@ class Model:
body_gravcomp: antigravity force, units of body weight (*, nbody)
body_contype: OR over all geom contypes (nbody,)
body_conaffinity: OR over all geom conaffinities (nbody,)
oct_child: octree children (noct, 8)
oct_aabb: octree axis-aligned bounding boxes (noct, 2, 3)
oct_coeff: octree interpolation coefficients (noct, 8)
@@ -1179,7 +1201,6 @@ class Model:
flex_stiffness: finite element stiffness matrix (nflexstiffness,)
flex_bending: bending stiffness (nflexbending,)
flex_damping: Rayleigh's damping coefficient (nflex,)
flex_edgeequality: edge equality type (0:none,1:edge,2:vert,3:strain) (nflex,)
flex_centered: flex vertices are centered at body origin (nflex,)
flexedge_J_rownnz: number of nonzeros in Jacobian row (nflexedge,)
@@ -1329,15 +1350,9 @@ class Model:
nmaxpolygon: maximum number of verts per polygon
nmaxmeshdeg: maximum number of polygons per vert
is_sparse: constraint Jacobian/Hessian layout (sparse vs dense). Does not affect M, whose
factorization is a per-block decision -- see qLD_* and m_block_layout
qLD_has_dense: any M block factors as a packed dense block
qLD_has_simple: any M block is simple (diagonal -> 1/diag, no factorization)
qLD_has_sparse: any M block factors via sparse LDL (oversized block / tendon armature)
factorization is a per-block decision -- see M_tiles and m_block_layout
qLD_block_total: packed length of the dense region per world (also the offset of the LDL region)
qLD_block_adr: packed offset of each dof's diagonal block; 0 if sparse (nv,)
qLD_dof_dense: per-dof flag, 1 if the dof's block is dense (packed) (nv,)
qLD_dof_simple: per-dof flag, 1 if the dof's block is simple (diagonal) (nv,)
qLD_simple_dofs: indices of the simple (diagonal) dofs (nsimple,)
qLD_block_adr: packed factor offset; Q_LD_BLOCK_* sentinel otherwise (nv,)
has_fluid: True if wind, density, or viscosity are non-zero at put_model time
has_sdf_geom: whether the model contains SDF geoms
has_flex_selfcollide: whether any flex has self-collision enabled
@@ -1378,7 +1393,6 @@ class Model:
tendon_geom_adr: geom tendon address
tendon_limited_adr: addresses for limited tendons
max_ten_J_rownnz: maximum number of non-zeros in a tendon row
jtcj_max_pairs: bound on a contact's support-pair count, sizes the elliptic-cone JTCJ launch
ten_wrapadr_site: wrap object starting address for sites
ten_wrapnum_site: number of site wrap objects per tendon
wrap_jnt_adr: addresses for joint tendon wrap object
@@ -1413,8 +1427,8 @@ class Model:
sensor_rangefinder_bodyid: bodyid for rangefinder (nrangefinder,)
taxel_vertadr: tactile sensor vertex address (nsensortaxel,)
taxel_sensorid: address for tactile sensors
M_tiles: tiling configuration
qLD_updates: tuple of index triples for sparse factorization
M_tiles: scalar and tiled block-factorization groups
qLD_updates: sparse factor updates grouped by tree level
qLD_all_updates: tuple of all levels concatenated
qLD_level_offsets: tuple of start offsets for each level
M_fullm_i: sparse mass matrix addressing
@@ -1436,12 +1450,19 @@ class Model:
flex_evpairflexid: maps each element-vertex pair directly to its flexid (nflexevpair,)
flex_vertflexid: maps each vertex index directly to its flexid (nflexvert,)
flex_shelladr: maps each flex to its start shell index (nflex,)
flex_faceadr: maps each flex to its start face index (nflex,)
flex_cell_map: precomputed flex cell mapping (nflexintcell,)
flexstrain_J_rownnz: number of nonzeros in flex strain Jacobian row (neq_flexstrain,)
flexstrain_J_rownnz: number of nonzeros in flex strain Jacobian row (neq_flexstrain,)
flexstrain_J_rowadr: row start address in colind array (neq_flexstrain,)
flexstrain_J_colind: column indices in sparse flex strain Jacobian (nJfs,)
flexstrain_J_colind: column indices in sparse flex strain Jacobian (nJfs,)
neq_flexstrain: number of flex strain equality constraints
nJfs: number of non-zeros in sparse flex strain Jacobian
nflexbend_interp: number of interpolated bending edges
flex_bend_interp_map: mapping of interpolated bending edges to flex and (nflexbend_interp, 2)
local edge indices
nflexface: number of interpolated flex shell faces
flex_face_map: mapping of face index to flex and local element face indices
flex_face: global node indices of each face (nflexface, 9)
"""
nq: int
@@ -1510,6 +1531,7 @@ class Model:
body_treeid: array("nbody", int)
body_geomnum: array("nbody", int)
body_geomadr: array("nbody", int)
body_simple: array("nbody", int)
body_pos: array("*", "nbody", wp.vec3)
body_quat: array("*", "nbody", wp.quat)
body_ipos: array("*", "nbody", wp.vec3)
@@ -1753,9 +1775,9 @@ class Model:
actuator_actnum: array("nu", int)
actuator_trnid: array("nu", wp.vec2i)
actuator_cranklength: array("*", "nu", float)
actuator_dynprm: array("*", "nu", vec10f)
actuator_gainprm: array("*", "nu", vec10f)
actuator_biasprm: array("*", "nu", vec10f)
actuator_dynprm: array("*", "nu", vec10)
actuator_gainprm: array("*", "nu", vec10)
actuator_biasprm: array("*", "nu", vec10)
actuator_actlimited: array("nu", bool)
actuator_actrange: array("*", "nu", wp.vec2)
actuator_actearly: array("nu", bool)
@@ -1810,14 +1832,8 @@ class Model:
nmaxpolygon: int
nmaxmeshdeg: int
is_sparse: bool
qLD_has_dense: bool
qLD_has_simple: bool
qLD_has_sparse: bool
qLD_block_total: int
qLD_block_adr: wp.array[int]
qLD_dof_dense: wp.array[int]
qLD_dof_simple: wp.array[int]
qLD_simple_dofs: wp.array[int]
qLD_block_adr: array("nv", int)
has_fluid: bool
has_sdf_geom: bool
has_flex_selfcollide: bool
@@ -1825,97 +1841,102 @@ class Model:
has_3d_flex: bool
max_flex_dim: int
block_dim: BlockDim
body_tree: tuple[wp.array[int], ...]
body_branches: wp.array[int]
body_branch_start: wp.array[int]
body_tree: tuple[array("nbody", int), ...]
body_branches: array("nbody_branches", int)
body_branch_start: array("nbranch_start", int)
mocap_bodyid: array("nmocap", int)
body_fluid_ellipsoid: array("nbody", bool)
body_fluid_ellipsoid_adr: wp.array[int]
body_fluid_box_adr: wp.array[int]
jnt_limited_slide_hinge_adr: wp.array[int]
jnt_limited_ball_adr: wp.array[int]
body_fluid_ellipsoid_adr: array("nbody_fluid_ellipsoid", int)
body_fluid_box_adr: array("nbody_fluid_box", int)
jnt_limited_slide_hinge_adr: array("njnt_limited_slide_hinge", int)
jnt_limited_ball_adr: array("njnt_limited_ball", int)
body_isdofancestor: array("nbody", "nv_pad", int)
dof_tri_row: wp.array[int]
dof_tri_col: wp.array[int]
nxn_geom_pair: wp.array[wp.vec2i]
nxn_geom_pair_filtered: wp.array[wp.vec2i]
nxn_pairid: wp.array[wp.vec2i]
nxn_pairid_filtered: wp.array[wp.vec2i]
dof_tri_row: array("ndof_tri", int)
dof_tri_col: array("ndof_tri", int)
nxn_geom_pair: array("nnxn_geom_pair", wp.vec2i)
nxn_geom_pair_filtered: array("nnxn_geom_pair_filtered", wp.vec2i)
nxn_pairid: array("nnxn_geom_pair", wp.vec2i)
nxn_pairid_filtered: array("nnxn_geom_pair_filtered", wp.vec2i)
geom_pair_type_count: tuple[int, ...]
geom_plugin_index: array("ngeom", int)
eq_connect_adr: wp.array[int]
eq_wld_adr: wp.array[int]
eq_jnt_adr: wp.array[int]
eq_ten_adr: wp.array[int]
eq_flex_adr: wp.array[int]
eq_flexstrain_adr: wp.array[int]
tendon_jnt_adr: wp.array[int]
tendon_site_pair_adr: wp.array[int]
tendon_geom_adr: wp.array[int]
tendon_limited_adr: wp.array[int]
eq_connect_adr: array("neq_connect", int)
eq_wld_adr: array("neq_wld", int)
eq_jnt_adr: array("neq_jnt", int)
eq_ten_adr: array("neq_ten", int)
eq_flex_adr: array("neq_flex", int)
eq_flexstrain_adr: array("neq_flexstrain", int)
tendon_jnt_adr: array("ntendon_jnt", int)
tendon_site_pair_adr: array("ntendon_site_pair", int)
tendon_geom_adr: array("ntendon_geom", int)
tendon_limited_adr: array("ntendon_limited", int)
max_ten_J_rownnz: int
jtcj_max_pairs: int
ten_wrapadr_site: wp.array[int]
ten_wrapnum_site: wp.array[int]
wrap_jnt_adr: wp.array[int]
wrap_site_adr: wp.array[int]
wrap_site_pair_adr: wp.array[int]
wrap_geom_adr: wp.array[int]
ten_wrapadr_site: array("nten_wrapadr_site", int)
ten_wrapnum_site: array("ntendon", int)
wrap_jnt_adr: array("nwrap_jnt", int)
wrap_site_adr: array("nwrap_site", int)
wrap_site_pair_adr: array("nwrap_site_pair", int)
wrap_geom_adr: array("nwrap_geom", int)
wrap_pulley_scale: array("nwrap", float)
actuator_trntype_body_adr: wp.array[int]
sensor_pos_adr: wp.array[int]
sensor_limitpos_adr: wp.array[int]
sensor_vel_adr: wp.array[int]
sensor_limitvel_adr: wp.array[int]
sensor_acc_adr: wp.array[int]
sensor_rangefinder_adr: wp.array[int]
rangefinder_sensor_adr: wp.array[int]
sensor_collision_start_adr: wp.array[int]
actuator_trntype_body_adr: array("nacttrnbody", int)
sensor_pos_adr: array("nsensor_pos", int)
sensor_limitpos_adr: array("nsensor_limitpos", int)
sensor_vel_adr: array("nsensor_vel", int)
sensor_limitvel_adr: array("nsensor_limitvel", int)
sensor_acc_adr: array("nsensor_acc", int)
sensor_rangefinder_adr: array("nrangefinder", int)
rangefinder_sensor_adr: array("nsensor", int)
sensor_collision_start_adr: array("nsensor_collision_start_adr", int)
collision_sensor_adr: array("nsensor", int)
sensor_touch_adr: wp.array[int]
sensor_limitfrc_adr: wp.array[int]
sensor_touch_adr: array("nsensor_touch", int)
sensor_limitfrc_adr: array("nsensor_limitfrc", int)
sensor_e_potential: bool
sensor_e_kinetic: bool
sensor_tendonactfrc_adr: wp.array[int]
sensor_tendonactfrc_adr: array("nsensor_tendonactfrc", int)
sensor_subtree_vel: bool
sensor_contact_adr: array("nsensorcontact", int)
sensor_adr_to_contact_adr: array("nsensor", int)
sensor_rne_postconstraint: bool
sensor_rangefinder_bodyid: array("nrangefinder", int)
taxel_vertadr: array("nsensortaxel", int)
taxel_sensorid: wp.array[int]
taxel_sensorid: array("nsensortaxel", int)
M_tiles: tuple[TileSet, ...]
qLD_updates: tuple[wp.array[wp.vec3i], ...]
qLD_all_updates: wp.array[wp.vec3i]
qLD_level_offsets: wp.array[int]
qLD_updates: tuple[array("nqLD_all_updates", wp.vec3i), ...]
qLD_all_updates: array("nqLD_all_updates", wp.vec3i)
qLD_level_offsets: array("nqLD_level_offsets", int)
# TODO(team): Remove M_fullm_i/j and M_elemid by iterating the M CSR layout
# directly in the solver/derivative kernels
M_fullm_i: wp.array[int]
M_fullm_j: wp.array[int]
M_elemid: wp.array2d[int] # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
M_hinit_i: wp.array[int] # row index of each CSR M entry (for densifying M into the dense Newton H)
M_fullm_upper_i: wp.array[int]
M_fullm_upper_j: wp.array[int]
M_fullm_upper_elemid: wp.array[int]
qD_fullm_i: wp.array[int] # D-structure (full square) row indices for RNE derivatives
qD_fullm_j: wp.array[int] # D-structure (full square) column indices for RNE derivatives
M_fullm_i: array("nM_fullm", int)
M_fullm_j: array("nM_fullm", int)
M_elemid: array("nv", "nv", int) # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
M_hinit_i: array("nC", int) # row index of each CSR M entry (for densifying M into the dense Newton H)
M_fullm_upper_i: array("nM_fullm_upper", int)
M_fullm_upper_j: array("nM_fullm_upper", int)
M_fullm_upper_elemid: array("nM_fullm_upper", int)
qD_fullm_i: array("nqD_fullm", int) # D-structure (full square) row indices for RNE derivatives
qD_fullm_j: array("nqD_fullm", int) # D-structure (full square) column indices for RNE derivatives
# Gather-based sparse mul_m indices (thread per DOF, no atomics)
M_mulm_rowadr: wp.array[int] # start address for each row [nv+1]
M_mulm_col: wp.array[int] # column index to gather from
M_mulm_madr: wp.array[int] # matrix address to read
flexelem_geom_pair_filtered: wp.array[wp.vec2i]
flexvert_geom_pair_filtered: wp.array[wp.vec2i]
M_mulm_rowadr: array("nv_plus_1", int) # start address for each row [nv+1]
M_mulm_col: array("nM_mulm", int) # column index to gather from
M_mulm_madr: array("nM_mulm", int) # matrix address to read
flexelem_geom_pair_filtered: array("nflexelem_geom_pair_filtered", wp.vec2i)
flexvert_geom_pair_filtered: array("nflexvert_geom_pair_filtered", wp.vec2i)
flex_elemflexid: array("nflexelem", int)
flex_shellflexid: array("nflexshelldata", int)
flex_evpairflexid: array("nflexevpair", int)
flex_vertflexid: array("nflexvert", int)
flex_shelladr: array("nflex", int)
flex_faceadr: array("nflex", int)
flex_cell_map: array("nflexintcell", wp.vec4i)
flexstrain_J_rownnz: array("neq_flexstrain", int)
flexstrain_J_rowadr: array("neq_flexstrain", int)
flexstrain_J_colind: array("nJfs", int)
neq_flexstrain: int
nJfs: int
nflexbend_interp: int
flex_bend_interp_map: array("nflexbend_interp", wp.vec2i)
nflexface: int
flex_face_map: array("nflexface", wp.vec2i)
flex_face: array("nflexface", 9, int)
class ContactType(enum.IntFlag):
@@ -1982,9 +2003,9 @@ class Constraint:
Attributes:
type: constraint type (ConstraintType) (nworld, njmax)
id: id of object of specific type (nworld, njmax)
jtdaj_adr: first efc row of each JTDAJ block (nworld, njmax)
jtdaj_nrow: efc rows per JTDAJ block (nworld, njmax)
jtdaj_nblock: number of JTDAJ blocks (nworld,)
jtdaj_adr: first efc row of each JTDAJ block (nworld, njmax)
jtdaj_nrow: efc rows per JTDAJ block (nworld, njmax)
jtdaj_nblock: number of JTDAJ blocks (nworld,)
J_rownnz: number of non-zeros in J row (nworld, 0) dense
(nworld, njmax) sparse
J_rowadr: row start address in colind array (nworld, 0) dense
@@ -2014,8 +2035,8 @@ class Constraint:
jtdaj_nblock: array("nworld", int)
J_rownnz: array("nworld", "njmax", int)
J_rowadr: array("nworld", "njmax", int)
J_colind: wp.array3d[int]
J: wp.array3d[float]
J_colind: array("nworld", 1, "njmax_nnz", int)
J: array("nworld", 1, "njmax_nnz", float)
pos: array("nworld", "njmax", float)
margin: array("nworld", "njmax", float)
D: array("nworld", "njmax_pad", float)
@@ -2098,7 +2119,7 @@ class Data:
M: total inertia, CSR (nworld, nC)
qLD: per-block factor: packed dense region, then the nC (nworld, qLD_block_total + nC)
L'*D*L region at offset qLD_block_total (nC=0 if no sparse block)
qLDiagInv: 1/diag(D) for the sparse LDL region (nworld, nv)
qLDiagInv: reciprocal diagonal for compact and sparse blocks (nworld, nv)
tree_awake: is tree awake; 0: asleep; 1: awake (nworld, ntree)
body_awake: body sleep state (SleepState) (nworld, nbody)
body_awake_ind: indices of awake/static bodies (nworld, nbody)
@@ -2133,7 +2154,7 @@ class Data:
tree_island: island ID per tree (-1 if unconstrained) (nworld, ntree)
dof_island: island ID per DOF (-1 if unconstrained) (nworld, nv)
island_dofadr: island start address in dof vector (nworld, ntree)
island_idofadr: island start address in idof vector (nworld, ntree)
island_idofadr: island start address in idof vector (nworld, ntree)
island_nv: DOFs per island (nworld, ntree)
island_nefc: constraints per island (nworld, ntree)
island_ne: equality constraints per island (nworld, ntree)
@@ -2150,8 +2171,8 @@ class Data:
cdof_dof: compacted DOF -> global DOF; -1 if unused (nworld, nvmax_pad)
ctol: compacted-solve main tolerance (nv/nvmax_pad scaled) (1,)
cls_tol: compacted-solve linesearch tolerance (1,)
cdof_tri_row: row index of compacted Hessian dof-pairs (nvmax_pad^2,)
cdof_tri_col: col index of compacted Hessian dof-pairs (nvmax_pad^2,)
cdof_tri_row: row index of compacted Hessian dof-pairs (nvmax_pad_sq,)
cdof_tri_col: col index of compacted Hessian dof-pairs (nvmax_pad_sq,)
cM: compacted dense inertia (nworld, nvmax_pad, nvmax_pad)
cqLD: compacted upper Cholesky factor (nworld, nvmax_pad, nvmax_pad)
crhs: compacted smooth-solve right-hand side (nworld, nvmax_pad, 1)
@@ -2179,6 +2200,8 @@ class Data:
flex_aabb_max: dynamic flex object bounding box max (nworld, nflex, 3)
flexnode_xpos: cartesian flex node positions (nworld, nflexnode, 3)
overflow: overflow bitmask (OverflowType) (nworld,)
face_xpos: cartesian flex face positions (nworld, nflexface, 9, 3)
face_quat: cartesian flex face orientations (nworld, nflexface, 4)
"""
solver_niter: array("nworld", int)
@@ -2242,8 +2265,8 @@ class Data:
moment_colind: array("nworld", "nJmom", int)
actuator_moment: array("nworld", "nJmom", float)
crb: array("nworld", "nbody", vec10)
M: wp.array2d[float]
qLD: wp.array2d[float]
M: array("nworld", "nC", float)
qLD: array("nworld", "qld_total", float)
qLDiagInv: array("nworld", "nv", float)
tree_awake: array("nworld", "ntree", int)
body_awake: array("nworld", "nbody", int)
@@ -2292,21 +2315,21 @@ class Data:
ncdof: array("nworld", int)
dof_cdof: array("nworld", "nv", int)
cdof_dof: array("nworld", "nvmax_pad", int)
ctol: wp.array[float]
cls_tol: wp.array[float]
cdof_tri_row: wp.array[int]
cdof_tri_col: wp.array[int]
cM: wp.array3d[float]
cqLD: wp.array3d[float]
crhs: wp.array3d[float]
cx: wp.array3d[float]
cJ: wp.array3d[float]
cMa: wp.array2d[float]
cqfrc_smooth: wp.array2d[float]
cqacc_smooth: wp.array2d[float]
cqacc_warmstart: wp.array2d[float]
cqacc: wp.array2d[float]
cqfrc_constraint: wp.array2d[float]
ctol: array(1, float)
cls_tol: array(1, float)
cdof_tri_row: array("nvmax_pad_sq", int)
cdof_tri_col: array("nvmax_pad_sq", int)
cM: array("nworld", "nvmax_pad", "nvmax_pad", float)
cqLD: array("nworld", "nvmax_pad", "nvmax_pad", float)
crhs: array("nworld", "nvmax_pad", 1, float)
cx: array("nworld", "nvmax_pad", 1, float)
cJ: array("nworld", "njmax_pad", "nvmax_pad", float)
cMa: array("nworld", "nvmax_pad", float)
cqfrc_smooth: array("nworld", "nvmax_pad", float)
cqacc_smooth: array("nworld", "nvmax_pad", float)
cqacc_warmstart: array("nworld", "nvmax_pad", float)
cqacc: array("nworld", "nvmax_pad", float)
cqfrc_constraint: array("nworld", "nvmax_pad", float)
# warp only fields:
nworld: int
@@ -2323,6 +2346,8 @@ class Data:
flex_aabb_max: array("nworld", "nflex", wp.vec3)
flexnode_xpos: array("nworld", "nflexnode", wp.vec3)
overflow: array("nworld", int)
face_xpos: array("nworld", "nflexface", 9, wp.vec3)
face_quat: array("nworld", "nflexface", wp.quat)
@dataclasses.dataclass
@@ -2332,8 +2357,13 @@ class InverseContext:
Jaref: wp.array2d[float]
search_dot: wp.array[float]
done: wp.array[bool]
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
quad_changed_ids: wp.array2d[int]
quad_changed_count: wp.array[int]
state_changed_count: wp.array[int]
ls_exhausted: wp.array[bool]
# the full-coordinate Data, set by solve_compact (None natively)
compact_m_full: Optional["Model"] = None
compact_d_full: Optional["Data"] = None
@dataclasses.dataclass
@@ -2345,6 +2375,7 @@ class SolverContext:
done: wp.array[bool]
grad: wp.array2d[float]
grad_dot: wp.array[float]
newton_decrement: wp.array[float]
Mgrad: wp.array2d[float]
search: wp.array2d[float]
mv: wp.array2d[float]
@@ -2352,16 +2383,21 @@ class SolverContext:
quad: wp.array2d[wp.vec3]
alpha: wp.array[float]
grad_scale: wp.array[float]
state_changed_count: wp.array[int]
improvement: wp.array[float]
ls_exhausted: wp.array[bool]
search_unchanged: wp.array[bool]
prev_grad: wp.array2d[float]
prev_Mgrad: wp.array2d[float]
beta: wp.array[float]
beta_den: wp.array[float]
h: wp.array3d[float]
hfactor: wp.array3d[float]
# Incremental Hessian update (Newton only)
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
quad_changed_ids: wp.array2d[int]
quad_changed_count: wp.array[int]
# the full-coordinate Data, set by solve_compact (None natively)
compact_m_full: Optional["Model"] = None
compact_d_full: Optional["Data"] = None
@dataclasses.dataclass
@@ -2373,6 +2409,7 @@ class RenderContext:
cam_res: camera resolution for actively rendering cameras
cam_id_map: camera id map
use_textures: whether to use textures
use_fast_math: whether to enable fast math for the render kernel
use_shadows: whether to use shadows
use_ambient_lighting: top-level switch for ambient contributions
background_color: color used for missed rays when no skybox is rendered
@@ -2460,6 +2497,7 @@ class RenderContext:
cam_res: array("ncam", wp.vec2i)
cam_id_map: array("ncam", int)
use_textures: bool
use_fast_math: bool
use_shadows: bool
use_ambient_lighting: bool
background_color: wp.uint32
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+5 -5
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name="mujoco-warp"
version = "3.10.0.1"
version = "3.11.0"
# TODO(team): create a distribution list
authors = [
{name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -29,9 +29,9 @@ requires-python = ">=3.10"
dependencies = [
"absl-py",
"etils[epath]",
"mujoco>=3.9.0",
"mujoco>=3.11.0",
"numpy",
"warp-lang>=1.14",
"warp-lang>=1.15",
]
[[tool.uv.index]]
@@ -56,8 +56,8 @@ dev = [
"ruff",
"pygls>=1.0.0,<2.0.0",
"lsprotocol>=2023.0.1,<2024.0.0",
"mujoco>=3.8.0.dev0",
"warp-lang>=1.14",
"mujoco>=3.11.0",
"warp-lang>=1.15",
"mjviser>=0.0.10",
"pillow",
]
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+1 -1
View File
@@ -88,7 +88,7 @@ def _make_warp_step_fn(mjm, m, d, graph, ctrls=None):
graph = _compile_step(m, d) if wp.get_device().is_cuda else None
wp.copy(d.ctrl, wp.array([mjd.ctrl.astype(np.float32)]))
wp.copy(d.act, wp.array([mjd.act.astype(np.float32)]))
wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied.astype(np.float32)]))
wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied], dtype=wp.spatial_vector))
wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
mjx/mujoco/mjx/warp/bvh.py
+3 -1
View File
@@ -118,7 +118,9 @@ def _refit_bvh_jax_impl(
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([]),
stage_in_argnames=set(['geom_size', 'geom_xmat', 'geom_xpos']),
stage_in_argnames=set(
['flexvert_xpos', 'geom_size', 'geom_xmat', 'geom_xpos']
),
stage_out_argnames=set([]),
graph_mode=m.opt._impl.graph_mode,
has_side_effect=True,
mjx/mujoco/mjx/warp/collision_driver.py
+49 -2
View File
@@ -46,7 +46,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
@@ -395,6 +394,27 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'contact__worldid',
]),
stage_in_argnames=set([
'body_awake',
'contact__dim',
'contact__dist',
'contact__efc_address',
'contact__elem',
'contact__flex',
'contact__frame',
'contact__friction',
'contact__geom',
'contact__geomcollisionid',
'contact__includemargin',
'contact__pos',
'contact__solimp',
'contact__solref',
'contact__solreffriction',
'contact__type',
'contact__vert',
'contact__worldid',
'flex_aabb_max',
'flex_aabb_min',
'flexvert_xpos',
'geom_aabb',
'geom_friction',
'geom_gap',
@@ -407,6 +427,10 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'geom_xmat',
'geom_xpos',
'hfield_data',
'nacon',
'ncollision',
'opt__ccd_tolerance',
'overflow',
'pair_friction',
'pair_gap',
'pair_margin',
@@ -414,7 +438,30 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'pair_solref',
'pair_solreffriction',
]),
stage_out_argnames=set([]),
stage_out_argnames=set([
'contact__dim',
'contact__dist',
'contact__efc_address',
'contact__elem',
'contact__flex',
'contact__frame',
'contact__friction',
'contact__geom',
'contact__geomcollisionid',
'contact__includemargin',
'contact__pos',
'contact__solimp',
'contact__solref',
'contact__solreffriction',
'contact__type',
'contact__vert',
'contact__worldid',
'flex_aabb_max',
'flex_aabb_min',
'nacon',
'ncollision',
'overflow',
]),
graph_mode=m.opt._impl.graph_mode,
has_side_effect=False,
)
mjx/mujoco/mjx/warp/forward.py
+823 -308
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/render.py
+4
View File
@@ -163,6 +163,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
'cam_intrinsic',
'cam_xmat',
'cam_xpos',
'flexvert_xpos',
'geom_matid',
'geom_rgba',
'geom_size',
@@ -177,11 +178,14 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
'light_exponent',
'light_specular',
'light_type',
'light_xdir',
'light_xpos',
'mat_emission',
'mat_rgba',
'mat_shininess',
'mat_specular',
'mat_texid',
'mat_texrepeat',
]),
stage_out_argnames=set([]),
graph_mode=m.opt._impl.graph_mode,
mjx/mujoco/mjx/warp/smooth.py
+15 -2
View File
@@ -414,9 +414,21 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
'qpos',
'site_xpos',
'subtree_com',
'ten_J',
'ten_length',
'ten_wrapadr',
'ten_wrapnum',
'wrap_obj',
'wrap_xpos',
]),
stage_out_argnames=set([
'ten_J',
'ten_length',
'ten_wrapadr',
'ten_wrapnum',
'wrap_obj',
'wrap_xpos',
]),
stage_out_argnames=set(['ten_length']),
graph_mode=m.opt._impl.graph_mode,
has_side_effect=False,
)
@@ -555,6 +567,7 @@ def _com_pos_jax_impl(m: types.Model, d: types.Data):
'body_mass',
'body_subtreemass',
'cdof',
'cinert',
'subtree_com',
'xanchor',
'xaxis',
@@ -562,7 +575,7 @@ def _com_pos_jax_impl(m: types.Model, d: types.Data):
'xipos',
'xmat',
]),
stage_out_argnames=set(['cdof', 'subtree_com']),
stage_out_argnames=set(['cdof', 'cinert', 'subtree_com']),
graph_mode=m.opt._impl.graph_mode,
has_side_effect=False,
)
mjx/mujoco/mjx/warp/types.py
+46 -46
View File
@@ -53,12 +53,14 @@ class TileSet:
Attributes:
adr: address of each tile in the set
size: size of all the tiles in this set
elemid: flat per-block gather indices into CSR M for tile_load_indexed
(absent -> nC sentinel)
elemid: flat CSR gather indices for tiled blocks; empty selects the
native-layout scalar path
"""
adr: np.ndarray
size: int
elemid: typing.Optional[np.ndarray] = None
elemid: np.ndarray = dataclasses.field(
default_factory=lambda: np.array([], dtype=np.int32)
)
def __eq__(self, other) -> bool:
if self.__class__ is not other.__class__:
@@ -99,6 +101,7 @@ class BlockDim:
cholesky_factorize_solve: block-dense Cholesky factorize+solve block
dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
small_cholesky: scalar small-block Cholesky block dimension (smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block
@@ -125,6 +128,7 @@ class BlockDim:
cholesky_factorize: int = 32
cholesky_factorize_solve: int = 32
cholesky_solve: int = 64
small_cholesky: int = 64
solve_LD_sparse_fused: int = 128
update_gradient_cholesky: int = 64
update_gradient_cholesky_blocked: int = 32
@@ -210,6 +214,7 @@ class ModelWarp(PyTreeNode):
eq_jnt_adr: np.ndarray
eq_ten_adr: np.ndarray
eq_wld_adr: np.ndarray
flex_bend_interp_map: np.ndarray
flex_bending: np.ndarray
flex_bendingadr: np.ndarray
flex_cell_map: np.ndarray
@@ -236,6 +241,9 @@ class ModelWarp(PyTreeNode):
flex_evpairadr: np.ndarray
flex_evpairflexid: np.ndarray
flex_evpairnum: np.ndarray
flex_face: np.ndarray
flex_face_map: np.ndarray
flex_faceadr: np.ndarray
flex_friction: np.ndarray
flex_gap: np.ndarray
flex_internal: np.ndarray
@@ -277,7 +285,6 @@ class ModelWarp(PyTreeNode):
is_sparse: bool
jnt_limited_ball_adr: np.ndarray
jnt_limited_slide_hinge_adr: np.ndarray
jtcj_max_pairs: int
light_bodyid: np.ndarray
light_targetbodyid: np.ndarray
mapD2M: np.ndarray
@@ -301,12 +308,14 @@ class ModelWarp(PyTreeNode):
nacttrnbody: int
nbranch: int
neq_flexstrain: int
nflexbend_interp: int
nflexbending: int
nflexedge: int
nflexelem: int
nflexelemdata: int
nflexelemedge: int
nflexevpair: int
nflexface: int
nflexintcell: int
nflexnode: int
nflexshelldata: int
@@ -338,13 +347,7 @@ class ModelWarp(PyTreeNode):
qLD_all_updates: np.ndarray
qLD_block_adr: np.ndarray
qLD_block_total: int
qLD_dof_dense: np.ndarray
qLD_dof_simple: np.ndarray
qLD_has_dense: bool
qLD_has_simple: bool
qLD_has_sparse: bool
qLD_level_offsets: np.ndarray
qLD_simple_dofs: np.ndarray
qLD_updates: Tuple[np.ndarray, ...]
rangefinder_sensor_adr: np.ndarray
sensor_acc_adr: np.ndarray
@@ -460,6 +463,8 @@ class DataWarp(PyTreeNode):
efc__vel: jax.Array
efc_islandid: jax.Array
energy: jax.Array
face_quat: jax.Array
face_xpos: jax.Array
flex_aabb_max: jax.Array
flex_aabb_min: jax.Array
flexedge_J: jax.Array
@@ -523,9 +528,6 @@ class DataWarp(PyTreeNode):
wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
'cJ',
'cM',
'cMa',
'cdof_tri_col',
'cdof_tri_row',
'cls_tol',
@@ -546,15 +548,7 @@ DATA_NON_VMAP = {
'contact__type',
'contact__vert',
'contact__worldid',
'cqLD',
'cqacc',
'cqacc_smooth',
'cqacc_warmstart',
'cqfrc_constraint',
'cqfrc_smooth',
'crhs',
'ctol',
'cx',
'naccdmax',
'nacon',
'naconmax',
@@ -673,6 +667,8 @@ _NDIM = {
'efc_islandid': 2,
'energy': 2,
'eq_active': 2,
'face_quat': 3,
'face_xpos': 4,
'flex_aabb_max': 3,
'flex_aabb_min': 3,
'flexedge_J': 2,
@@ -826,6 +822,7 @@ _NDIM = {
'block_dim__ray': 0,
'block_dim__render': 0,
'block_dim__segmented_sort': 0,
'block_dim__small_cholesky': 0,
'block_dim__solve_LD_sparse_fused': 0,
'block_dim__solve_beta_accumulate': 0,
'block_dim__solve_init_search_cg': 0,
@@ -860,6 +857,7 @@ _NDIM = {
'body_pos': 3,
'body_quat': 3,
'body_rootid': 1,
'body_simple': 1,
'body_subtreemass': 2,
'body_tree': -1,
'body_treeid': 1,
@@ -908,6 +906,7 @@ _NDIM = {
'eq_type': 1,
'eq_wld_adr': 1,
'exclude_signature': 1,
'flex_bend_interp_map': 2,
'flex_bending': 1,
'flex_bendingadr': 1,
'flex_cell_map': 2,
@@ -934,6 +933,9 @@ _NDIM = {
'flex_evpairadr': 1,
'flex_evpairflexid': 1,
'flex_evpairnum': 1,
'flex_face': 2,
'flex_face_map': 2,
'flex_faceadr': 1,
'flex_friction': 2,
'flex_gap': 1,
'flex_internal': 1,
@@ -1026,7 +1028,6 @@ _NDIM = {
'jnt_stiffness': 2,
'jnt_stiffnesspoly': 3,
'jnt_type': 1,
'jtcj_max_pairs': 0,
'light_active': 2,
'light_ambient': 3,
'light_attenuation': 3,
@@ -1094,12 +1095,14 @@ _NDIM = {
'neq_flexstrain': 0,
'nexclude': 0,
'nflex': 0,
'nflexbend_interp': 0,
'nflexbending': 0,
'nflexedge': 0,
'nflexelem': 0,
'nflexelemdata': 0,
'nflexelemedge': 0,
'nflexevpair': 0,
'nflexface': 0,
'nflexintcell': 0,
'nflexnode': 0,
'nflexshelldata': 0,
@@ -1191,13 +1194,7 @@ _NDIM = {
'qLD_all_updates': 2,
'qLD_block_adr': 1,
'qLD_block_total': 0,
'qLD_dof_dense': 1,
'qLD_dof_simple': 1,
'qLD_has_dense': 0,
'qLD_has_simple': 0,
'qLD_has_sparse': 0,
'qLD_level_offsets': 1,
'qLD_simple_dofs': 1,
'qLD_updates': -1,
'qpos0': 2,
'qpos_spring': 2,
@@ -1323,9 +1320,9 @@ _BATCH_DIM = {
'actuator_velocity': True,
'body_awake': True,
'body_awake_ind': True,
'cJ': False,
'cM': False,
'cMa': False,
'cJ': True,
'cM': True,
'cMa': True,
'cacc': True,
'cam_xmat': True,
'cam_xpos': True,
@@ -1355,18 +1352,18 @@ _BATCH_DIM = {
'contact__type': False,
'contact__vert': False,
'contact__worldid': False,
'cqLD': False,
'cqacc': False,
'cqacc_smooth': False,
'cqacc_warmstart': False,
'cqfrc_constraint': False,
'cqfrc_smooth': False,
'cqLD': True,
'cqacc': True,
'cqacc_smooth': True,
'cqacc_warmstart': True,
'cqfrc_constraint': True,
'cqfrc_smooth': True,
'crb': True,
'crhs': False,
'crhs': True,
'ctol': False,
'ctrl': True,
'cvel': True,
'cx': False,
'cx': True,
'dof_awake_ind': True,
'dof_cdof': True,
'dof_island': True,
@@ -1394,6 +1391,8 @@ _BATCH_DIM = {
'efc_islandid': True,
'energy': True,
'eq_active': True,
'face_quat': True,
'face_xpos': True,
'flex_aabb_max': True,
'flex_aabb_min': True,
'flexedge_J': True,
@@ -1547,6 +1546,7 @@ _BATCH_DIM = {
'block_dim__ray': False,
'block_dim__render': False,
'block_dim__segmented_sort': False,
'block_dim__small_cholesky': False,
'block_dim__solve_LD_sparse_fused': False,
'block_dim__solve_beta_accumulate': False,
'block_dim__solve_init_search_cg': False,
@@ -1581,6 +1581,7 @@ _BATCH_DIM = {
'body_pos': True,
'body_quat': True,
'body_rootid': False,
'body_simple': False,
'body_subtreemass': True,
'body_tree': False,
'body_treeid': False,
@@ -1629,6 +1630,7 @@ _BATCH_DIM = {
'eq_type': False,
'eq_wld_adr': False,
'exclude_signature': False,
'flex_bend_interp_map': False,
'flex_bending': False,
'flex_bendingadr': False,
'flex_cell_map': False,
@@ -1655,6 +1657,9 @@ _BATCH_DIM = {
'flex_evpairadr': False,
'flex_evpairflexid': False,
'flex_evpairnum': False,
'flex_face': False,
'flex_face_map': False,
'flex_faceadr': False,
'flex_friction': False,
'flex_gap': False,
'flex_internal': False,
@@ -1747,7 +1752,6 @@ _BATCH_DIM = {
'jnt_stiffness': True,
'jnt_stiffnesspoly': True,
'jnt_type': False,
'jtcj_max_pairs': False,
'light_active': True,
'light_ambient': True,
'light_attenuation': True,
@@ -1815,12 +1819,14 @@ _BATCH_DIM = {
'neq_flexstrain': False,
'nexclude': False,
'nflex': False,
'nflexbend_interp': False,
'nflexbending': False,
'nflexedge': False,
'nflexelem': False,
'nflexelemdata': False,
'nflexelemedge': False,
'nflexevpair': False,
'nflexface': False,
'nflexintcell': False,
'nflexnode': False,
'nflexshelldata': False,
@@ -1912,13 +1918,7 @@ _BATCH_DIM = {
'qLD_all_updates': False,
'qLD_block_adr': False,
'qLD_block_total': False,
'qLD_dof_dense': False,
'qLD_dof_simple': False,
'qLD_has_dense': False,
'qLD_has_simple': False,
'qLD_has_sparse': False,
'qLD_level_offsets': False,
'qLD_simple_dofs': False,
'qLD_updates': False,
'qpos0': True,
'qpos_spring': True,
mjx/mujoco/mjx/warp/types_test.py
+70 -1
View File
@@ -14,11 +14,80 @@
# ==============================================================================
"""Tests for generated MJX Warp types."""
import dataclasses
from absl.testing import absltest
from absl.testing import parameterized
import mujoco
import mujoco.mjx.warp as mjxw
from mujoco.mjx.warp import mjwp_io
from mujoco.mjx.warp import types
from mujoco.mjx.warp import warp as wp
import numpy as np
from mujoco.mjx.warp import types
class GeneratedTypesMetadataTest(absltest.TestCase):
def setUp(self):
super().setUp()
if not mjxw.WARP_INSTALLED or mjwp_io is None:
self.skipTest('warp is not installed')
def test_batched_fields_metadata(self):
m = mujoco.MjModel.from_xml_string('<mujoco/>')
d = mujoco.MjData(m)
mw = mjwp_io.put_model(m)
dw = mjwp_io.put_data(m, d)
def _check_fields(cls_name, obj, warp_cls, prefix=''):
for f in dataclasses.fields(obj):
val = getattr(obj, f.name)
if dataclasses.is_dataclass(val):
sub_cls_name = type(val).__name__
if sub_cls_name in types._BATCH_DIM and hasattr(
types, f'{sub_cls_name}Warp'
):
# Dedicated top-level sub-class (Option, Statistic)
_check_fields(
sub_cls_name, val, getattr(types, f'{sub_cls_name}Warp')
)
else:
# Flattened nested sub-dataclass (Data.contact, Data.efc)
_check_fields(cls_name, val, warp_cls, prefix=f'{prefix}{f.name}__')
continue
field_name = f'{prefix}{f.name}'
is_batched = getattr(val, '_is_batched', False)
if is_batched:
self.assertTrue(
types._BATCH_DIM[cls_name].get(field_name, False),
f'Expected {cls_name}.{field_name} with _is_batched=True to be'
' True in _BATCH_DIM',
)
if cls_name == 'Data':
self.assertNotIn(
field_name,
types.DATA_NON_VMAP,
f'Expected batched Data field {field_name} to not be in'
' DATA_NON_VMAP',
)
ann = warp_cls.__annotations__.get(field_name)
if ann is not None:
self.assertIn(
str(ann),
('jax.Array', "<class 'jax.Array'>"),
f'Expected {cls_name}.{field_name} to have jax.Array'
' annotation',
)
elif isinstance(val, wp.array):
self.assertFalse(
types._BATCH_DIM[cls_name].get(field_name, True),
f'Expected {cls_name}.{field_name} with _is_batched=False to be'
' False in _BATCH_DIM',
)
_check_fields('Model', mw, types.ModelWarp)
_check_fields('Data', dw, types.DataWarp)
class TileSetTest(parameterized.TestCase):