diff --git a/mjx/mujoco/mjx/codegen/generate_warp_shim.py b/mjx/mujoco/mjx/codegen/generate_warp_shim.py
index 46516f88..b562f0db 100644
--- a/mjx/mujoco/mjx/codegen/generate_warp_shim.py
+++ b/mjx/mujoco/mjx/codegen/generate_warp_shim.py
@@ -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)
diff --git a/mjx/mujoco/mjx/codegen/generate_warp_types.py b/mjx/mujoco/mjx/codegen/generate_warp_types.py
index 9f930611..33a0b583 100644
--- a/mjx/mujoco/mjx/codegen/generate_warp_types.py
+++ b/mjx/mujoco/mjx/codegen/generate_warp_types.py
@@ -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 = """
-
-
-
-
-
-
-
-
-
-
-
-
-"""
+
+_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))
diff --git a/mjx/mujoco/mjx/codegen/trace.py b/mjx/mujoco/mjx/codegen/trace.py
index f8339aa1..0207652f 100644
--- a/mjx/mujoco/mjx/codegen/trace.py
+++ b/mjx/mujoco/mjx/codegen/trace.py
@@ -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)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
index ebfc44a6..2d55eeae 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
index 7faaa578..190a147a 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
@@ -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)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
index 37823711..244cbc66 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
index 1accd276..8a9885b5 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
@@ -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],
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
index 0e906855..12798e8d 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
index e65bab44..a5880ed5 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
@@ -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],
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py
index 9e165b14..20b375d5 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py
@@ -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,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
index 7776ec99..0af054d4 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
@@ -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,
)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py
index 73946675..a4ec8f4e 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
index a96c8985..4feb854c 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
index 67f5c401..cc94c27e 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
@@ -28,7 +28,33 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec11
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
+def _add_weight(
+ # In:
+ nb: int,
+ body: types.vec16i,
+ weight: types.vec16,
+ b: int,
+ w: float,
+ check_weld: bool,
+) -> tuple[int, types.vec16i, types.vec16]:
+ """Appends weight to body weld-id, merging duplicates by summing weights."""
+ if wp.abs(w) < 1.0e-10:
+ return nb, body, weight
+ for i in range(16):
+ if i >= nb:
+ break
+ if body[i] == b:
+ weight[i] += w
+ return nb, body, weight
+ if nb < 16:
+ body[nb] = b
+ weight[nb] = w
+ return nb + 1, body, weight
+ return nb, body, weight
@wp.kernel
@@ -127,7 +153,7 @@ def _efc_row(
@cache_kernel
def _equality_connect(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -471,7 +497,7 @@ def _equality_connect(is_sparse: bool, newton: bool):
@cache_kernel
def _equality_joint(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -613,7 +639,7 @@ def _equality_joint(is_sparse: bool, newton: bool):
@cache_kernel
def _equality_tendon(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -802,7 +828,7 @@ def _equality_tendon(is_sparse: bool, newton: bool):
@cache_kernel
def _equality_flex(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -937,7 +963,7 @@ def _equality_flex(is_sparse: bool, newton: bool):
@cache_kernel
def _equality_weld(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -1072,7 +1098,7 @@ def _equality_weld(is_sparse: bool, newton: bool):
qfull1 = math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])
qdot1 = math.mul_quat(omega2_q, qfull1) * 0.5
- negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+ negqdot1 = math.quat_inv(qdot1)
negq1 = wp.quat(qfull1[0], -qfull1[1], -qfull1[2], -qfull1[3])
else:
@@ -1083,7 +1109,7 @@ def _equality_weld(is_sparse: bool, newton: bool):
q1_non_site = xquat_in[worldid, body2]
qdot1 = math.mul_quat(omega2_q, q1_non_site) * 0.5
- negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+ negqdot1 = math.quat_inv(qdot1)
negq1 = wp.quat(q1_non_site[0], -q1_non_site[1], -q1_non_site[2], -q1_non_site[3])
# compute Jacobian difference (opposite of contact: 0 - 1)
@@ -1414,7 +1440,7 @@ def _equality_weld(is_sparse: bool, newton: bool):
@cache_kernel
def _equality_flexstrain(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -1455,6 +1481,7 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
njmax_in: int,
njmax_nnz_in: int,
flexnode_xpos_in: wp.array2d[wp.vec3],
+ face_quat_in: wp.array2d[wp.quat],
# Data out:
ne_out: wp.array[int],
nefc_out: wp.array[int],
@@ -1484,11 +1511,11 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
f = eq_obj1id[eqid]
order = flex_interp[f]
- if order <= 0:
+ if order == 0:
return
# nodes per cell
- npc = (order + 1) * (order + 1) * (order + 1)
+ npc = wp.where(order < 0, 4, 8)
# cell indices from eq_data
data = eq_data[worldid % eq_data.shape[0], eqid]
@@ -1497,16 +1524,25 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
ck = int(data[2])
cellnum = flex_cellnum[f]
+ cx = cellnum[0]
cy = cellnum[1]
cz = cellnum[2]
nstart = flex_nodeadr[f]
- ny_g = cy * order + 1
- nz_g = cz * order + 1
+ ny_g = cy + 1
+ nz_g = cz + 1
ndof_cell = 3 * npc
# read eigenmode data from flex_stiffness
- cell_idx = ci * cy * cz + cj * cz + ck
+ cell_idx = wp.where(order < 0, ci, ci * cy * cz + cj * cz + ck)
+ normal_axis = 0
+ g_fixed = 0
+ q0 = 0
+ q1 = 0
+ ny_g_face = 0
+ nz_g_face = 0
+ if order < 0:
+ normal_axis, g_fixed, q0, q1, ny_g_face, nz_g_face = support.get_face_metadata(cx, cy, cz, cell_idx, 1)
k_base = flex_stiffnessadr[f] + cell_idx * ndof_cell * ndof_cell
neig = int(flex_stiffness[k_base])
@@ -1520,20 +1556,46 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
# We compute the corotational quaternion from the deformation gradient
# at the cell center (0.5, 0.5, 0.5)
- cell_quat = support.compute_interp_cell_quat(flexnode_xpos_in, order, ci, cj, ck, cy, cz, ny_g, nz_g, nstart, worldid)
+ cell_quat = wp.quat(0.0, 0.0, 0.0, 1.0)
+ if order < 0:
+ cell_quat = face_quat_in[worldid, cell_idx]
+ else:
+ cell_quat = support.compute_interp_cell_quat(
+ flexnode_xpos_in,
+ 1,
+ ci,
+ cj,
+ ck,
+ cy,
+ cz,
+ ny_g,
+ nz_g,
+ nstart,
+ worldid,
+ )
cell_quat_inv = wp.quat(-cell_quat[0], -cell_quat[1], -cell_quat[2], cell_quat[3])
# Compute average invweight across cell nodes (translation component)
avg_invweight = float(0.0)
idx_iw = int(0)
- for li_iw in range(order + 1):
- for lj_iw in range(order + 1):
- for lk_iw in range(order + 1):
+ for li_iw in range(2):
+ for lj_iw in range(2):
+ for lk_iw in range(2):
if idx_iw < npc:
- gi_iw = ci * order + li_iw
- gj_iw = cj * order + lj_iw
- gk_iw = ck * order + lk_iw
- gidx_iw = gi_iw * ny_g * nz_g + gj_iw * nz_g + gk_iw
+ gidx_iw = wp.where(
+ order < 0,
+ support.gather_face_node_index_fast(
+ normal_axis,
+ g_fixed,
+ q0,
+ q1,
+ ny_g_face,
+ nz_g_face,
+ idx_iw,
+ 1,
+ ),
+ (ci + li_iw) * ny_g * nz_g + (cj + lj_iw) * nz_g + (ck + lk_iw),
+ )
bodyid_iw = flex_nodebodyid[nstart + gidx_iw]
avg_invweight += body_invweight0[worldid % body_invweight0.shape[0], bodyid_iw][0]
idx_iw += 1
@@ -1558,14 +1620,24 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
# Compute constraint residual: dot(eigvec, displacement_in_corot_frame)
residual = float(0.0)
idx2 = int(0)
- for li2 in range(order + 1):
- for lj2 in range(order + 1):
- for lk2 in range(order + 1):
+ for li2 in range(2):
+ for lj2 in range(2):
+ for lk2 in range(2):
if idx2 < npc:
- gi2 = ci * order + li2
- gj2 = cj * order + lj2
- gk2 = ck * order + lk2
- gidx2 = gi2 * ny_g * nz_g + gj2 * nz_g + gk2
+ gidx2 = wp.where(
+ order < 0,
+ support.gather_face_node_index_fast(
+ normal_axis,
+ g_fixed,
+ q0,
+ q1,
+ ny_g_face,
+ nz_g_face,
+ idx2,
+ 1,
+ ),
+ (ci + li2) * ny_g * nz_g + (cj + lj2) * nz_g + (ck + lk2),
+ )
xpos_n = flexnode_xpos_in[worldid, nstart + gidx2]
refpos_n = flex_node0[nstart + gidx2]
@@ -1606,14 +1678,24 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
q = flexstrain_J_colind[fs_rowadr + sparseid]
J_val = float(0.0)
idx3 = int(0)
- for li3 in range(order + 1):
- for lj3 in range(order + 1):
- for lk3 in range(order + 1):
+ for li3 in range(2):
+ for lj3 in range(2):
+ for lk3 in range(2):
if idx3 < npc:
- gi3 = ci * order + li3
- gj3 = cj * order + lj3
- gk3 = ck * order + lk3
- gidx3 = gi3 * ny_g * nz_g + gj3 * nz_g + gk3
+ gidx3 = wp.where(
+ order < 0,
+ support.gather_face_node_index_fast(
+ normal_axis,
+ g_fixed,
+ q0,
+ q1,
+ ny_g_face,
+ nz_g_face,
+ idx3,
+ 1,
+ ),
+ (ci + li3) * ny_g * nz_g + (cj + lj3) * nz_g + (ck + lk3),
+ )
bodyid3 = flex_nodebodyid[nstart + gidx3]
xpos_n3 = flexnode_xpos_in[worldid, nstart + gidx3]
@@ -1681,7 +1763,7 @@ def _equality_flexstrain(is_sparse: bool, newton: bool):
@cache_kernel
def _friction_dof(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -1782,7 +1864,7 @@ def _friction_dof(is_sparse: bool, newton: bool):
@cache_kernel
def _friction_tendon(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -1906,7 +1988,7 @@ def _friction_tendon(is_sparse: bool, newton: bool):
@cache_kernel
def _limit_slide_hinge(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -2022,7 +2104,7 @@ def _limit_slide_hinge(is_sparse: bool, newton: bool):
@cache_kernel
def _limit_ball(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -2158,7 +2240,7 @@ def _limit_ball(is_sparse: bool, newton: bool):
@cache_kernel
def _limit_tendon(is_sparse: bool, newton: bool):
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
nv: int,
@@ -2560,7 +2642,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool, newton: bool):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
IS_SPARSE = is_sparse
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
def kernel(
# Model:
body_weldid: wp.array[int],
@@ -2674,7 +2756,7 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b
IS_SPARSE = is_sparse
HAS_FLEX = True
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
body_parentid: wp.array[int],
@@ -2831,108 +2913,118 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b
is_interp = True
if is_interp:
- # Interpolated flex: sum of all contributing body dofnums
- rownnz = int(0)
+ # Interpolated flex: sum of all contributing body dofnums after de-duplication
+ local_bodies = types.vec16i(-1)
+ local_weights = types.vec16(0.0)
+ local_nb = int(0)
+
for side in range(2):
if geom[side] >= 0:
b = body_weldid[geom_bodyid[geom[side]]]
- rownnz += body_dofnum[b]
- elif flex[side] >= 0 and vert[side] >= 0:
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, 1.0, True)
+ elif flex[side] >= 0:
f = flex[side]
if flex_interp[f] != 0:
- # Compute parametric coordinate from flex_vert0
- v_adr = flex_vertadr[f] + vert[side]
- coord = flex_vert0[v_adr]
- cn = flex_cellnum[f]
- cx = cn[0]
- cy = cn[1]
- cz = cn[2]
-
- # Cell lookup
- ci = wp.min(int(coord[0] * float(cx)), cx - 1)
- ci = wp.max(ci, 0)
- cj = wp.min(int(coord[1] * float(cy)), cy - 1)
- cj = wp.max(cj, 0)
- ck = wp.min(int(coord[2] * float(cz)), cz - 1)
- ck = wp.max(ck, 0)
-
- # Local parametric coordinates
- local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
- local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
- local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
- local = wp.vec3(local_x, local_y, local_z)
-
- # Node grid dimensions
- ny_g = cy + 1
- nz_g = cz + 1
- nstart = flex_nodeadr[f]
-
- # Loop over 8 trilinear nodes
- for li in range(2):
- for lj in range(2):
- for lk in range(2):
- w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
- if w > 1.0e-5:
- gi = ci + li
- gj = cj + lj
- gk = ck + lk
- node_idx = gi * ny_g * nz_g + gj * nz_g + gk
- b = body_weldid[flex_nodebodyid[nstart + node_idx]]
- rownnz += body_dofnum[b]
+ # Interpolated flex path (trilinear)
+ if vert[side] >= 0:
+ v_adr = flex_vertadr[f] + vert[side]
+ coord = flex_vert0[v_adr]
+ cn = flex_cellnum[f]
+ cx = cn[0]
+ cy = cn[1]
+ cz = cn[2]
+ ci = wp.min(int(coord[0] * float(cx)), cx - 1)
+ ci = wp.max(ci, 0)
+ cj = wp.min(int(coord[1] * float(cy)), cy - 1)
+ cj = wp.max(cj, 0)
+ ck = wp.min(int(coord[2] * float(cz)), cz - 1)
+ ck = wp.max(ck, 0)
+ local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
+ local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
+ local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
+ local = wp.vec3(local_x, local_y, local_z)
+ ny_g = cy + 1
+ nz_g = cz + 1
+ nstart = flex_nodeadr[f]
+ for li in range(2):
+ for lj in range(2):
+ for lk in range(2):
+ w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
+ if w > 1.0e-5:
+ gi = ci + li
+ gj = cj + lj
+ gk = ck + lk
+ node_idx = gi * ny_g * nz_g + gj * nz_g + gk
+ b = body_weldid[flex_nodebodyid[nstart + node_idx]]
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, w, True)
+ elif elem[side] >= 0:
+ e = elem[side]
+ dim_f = flex_dim[f]
+ edata_adr = flex_elemdataadr[f] + e * (dim_f + 1)
+ vert_adr_f = flex_vertadr[f]
+ contact_pos = pos_in[conid]
+ total_inv_dist = float(0.0)
+ blended_coord = wp.vec3(0.0, 0.0, 0.0)
+ for vi in range(4):
+ if vi <= dim_f:
+ v_idx = flex_elem[edata_adr + vi]
+ vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx]
+ dist_v = wp.length(contact_pos - vpos)
+ w_inv = 1.0 / wp.max(1.0e-10, dist_v)
+ total_inv_dist += w_inv
+ blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv
+ if total_inv_dist > 1.0e-10:
+ blended_coord = blended_coord / total_inv_dist
+ cn = flex_cellnum[f]
+ cx = cn[0]
+ cy = cn[1]
+ cz = cn[2]
+ ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1)
+ ci = wp.max(ci, 0)
+ cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1)
+ cj = wp.max(cj, 0)
+ ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1)
+ ck = wp.max(ck, 0)
+ local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0)
+ local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0)
+ local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0)
+ local = wp.vec3(local_x, local_y, local_z)
+ ny_g = cy + 1
+ nz_g = cz + 1
+ nstart = flex_nodeadr[f]
+ for li in range(2):
+ for lj in range(2):
+ for lk in range(2):
+ w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
+ if w > 1.0e-5:
+ gi = ci + li
+ gj = cj + lj
+ gk = ck + lk
+ node_idx = gi * ny_g * nz_g + gj * nz_g + gk
+ b = body_weldid[flex_nodebodyid[nstart + node_idx]]
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, w, True)
else:
- b = body_weldid[flex_vertbodyid[flex_vertadr[f] + vert[side]]]
- rownnz += body_dofnum[b]
- elif flex[side] >= 0 and elem[side] >= 0:
- # Elem contact: use blended coordinate from distance weighting
- f = flex[side]
- e = elem[side]
- dim_f = flex_dim[f]
- edata_adr = flex_elemdataadr[f] + e * (dim_f + 1)
- vert_adr_f = flex_vertadr[f]
- contact_pos = pos_in[conid]
+ # Non-interpolated flex path: use pre-computed body_ids and weights
+ if side == 0:
+ for vi in range(4):
+ if body_ids1[vi] >= 0:
+ b = body_weldid[body_ids1[vi]]
+ local_nb, local_bodies, local_weights = _add_weight(
+ local_nb, local_bodies, local_weights, b, weights1[vi], True
+ )
+ else:
+ for vi in range(4):
+ if body_ids2[vi] >= 0:
+ b = body_weldid[body_ids2[vi]]
+ local_nb, local_bodies, local_weights = _add_weight(
+ local_nb, local_bodies, local_weights, b, weights2[vi], True
+ )
- total_inv_dist = float(0.0)
- blended_coord = wp.vec3(0.0, 0.0, 0.0)
- for vi in range(4):
- if vi <= dim_f:
- v_idx = flex_elem[edata_adr + vi]
- vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx]
- dist_v = wp.length(contact_pos - vpos)
- w_inv = 1.0 / wp.max(1.0e-10, dist_v)
- total_inv_dist += w_inv
- blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv
- if total_inv_dist > 1.0e-10:
- blended_coord = blended_coord / total_inv_dist
-
- if flex_interp[f] != 0:
- cn = flex_cellnum[f]
- cx = cn[0]
- cy = cn[1]
- cz = cn[2]
- ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1)
- ci = wp.max(ci, 0)
- cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1)
- cj = wp.max(cj, 0)
- ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1)
- ck = wp.max(ck, 0)
- local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0)
- local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0)
- local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0)
- local = wp.vec3(local_x, local_y, local_z)
- ny_g = cy + 1
- nz_g = cz + 1
- nstart = flex_nodeadr[f]
- for li in range(2):
- for lj in range(2):
- for lk in range(2):
- w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
- if w > 1.0e-5:
- gi = ci + li
- gj = cj + lj
- gk = ck + lk
- node_idx = gi * ny_g * nz_g + gj * nz_g + gk
- b = body_weldid[flex_nodebodyid[nstart + node_idx]]
- rownnz += body_dofnum[b]
+ # sum dofnums for unique bodies
+ rownnz = int(0)
+ for i in range(16):
+ if i < local_nb:
+ rownnz += body_dofnum[local_bodies[i]]
else:
# Standard path (including elements up to 4 bodies)
b1_0 = body_weldid[body_ids1[0]]
@@ -3000,7 +3092,7 @@ def _efc_contact_init_flex(cone_type: types.ConeType, is_sparse: bool, newton: b
def _efc_contact_jac_sparse(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
body_parentid: wp.array[int],
@@ -3153,7 +3245,7 @@ def _efc_contact_jac_sparse_flex(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
HAS_FLEX = True
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
body_parentid: wp.array[int],
@@ -3238,308 +3330,208 @@ def _efc_contact_jac_sparse_flex(cone_type: types.ConeType):
rownnz = efc_J_rownnz_in[worldid, efcid]
if is_interp:
- # Interpolated flex path: iterate over bodies per side, accumulate weighted Jacobians
- nnz = int(0)
- Jqvel = float(0.0)
+ # Interpolated flex path: accumulate unique bodies and signed weights
+ local_bodies = types.vec16i(-1)
+ local_weights = types.vec16(0.0)
+ local_nb = int(0)
for side in range(2):
sign = float(-1.0) if side == 0 else float(1.0)
if geom[side] >= 0:
- # Geom side: single body
b = body_weldid[geom_bodyid[geom[side]]]
- dof_start = body_dofadr[b]
- ndof = body_dofnum[b]
- for di in range(ndof):
- dofid = dof_start + di
- jacp, jacr = support.jac_dof(
- body_parentid,
- body_rootid,
- dof_bodyid,
- body_isdofancestor,
- subtree_com_in,
- cdof_in,
- con_pos,
- b,
- dofid,
- worldid,
- )
-
- if wp.static(IS_ELLIPTIC):
- J = float(0.0)
- if dimid < 3:
- frame_row = frame_in[conid, dimid]
- for xyz in range(3):
- J += frame_row[xyz] * jacp[xyz] * sign
- else:
- frame_row = frame_in[conid, dimid - 3]
- for xyz in range(3):
- J += frame_row[xyz] * jacr[xyz] * sign
- else:
- J = float(0.0)
- Ji = float(0.0)
- for xyz in range(3):
- J += frame_0[xyz] * jacp[xyz] * sign
- if condim > 1:
- if dimid2 < 3:
- Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * sign
- else:
- Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * sign
- if condim > 1:
- if dimid % 2 == 0:
- J += Ji * frii
- else:
- J -= Ji * frii
-
- if nnz < rownnz:
- sparseid = rowadr + nnz
- efc_J_colind_out[worldid, 0, sparseid] = dofid
- efc_J_out[worldid, 0, sparseid] = J
- Jqvel += J * qvel_in[worldid, dofid]
- nnz += 1
-
- elif flex[side] >= 0 and vert[side] >= 0:
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, sign, True)
+ elif flex[side] >= 0:
f = flex[side]
if flex_interp[f] != 0:
- # Interpolated flex side: compute cell node bodies and weights
- v_adr = flex_vertadr[f] + vert[side]
- coord = flex_vert0[v_adr]
- cn = flex_cellnum[f]
- cx = cn[0]
- cy = cn[1]
- cz = cn[2]
-
- ci = wp.min(int(coord[0] * float(cx)), cx - 1)
- ci = wp.max(ci, 0)
- cj = wp.min(int(coord[1] * float(cy)), cy - 1)
- cj = wp.max(cj, 0)
- ck = wp.min(int(coord[2] * float(cz)), cz - 1)
- ck = wp.max(ck, 0)
-
- local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
- local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
- local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
- local = wp.vec3(local_x, local_y, local_z)
-
- ny_g = cy + 1
- nz_g = cz + 1
- nstart = flex_nodeadr[f]
-
- for li in range(2):
- for lj in range(2):
- for lk in range(2):
- w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
- if w > 1.0e-5:
- gi = ci + li
- gj = cj + lj
- gk = ck + lk
- node_idx = gi * ny_g * nz_g + gj * nz_g + gk
- b = body_weldid[flex_nodebodyid[nstart + node_idx]]
- w_sign = w * sign
-
- dof_start = body_dofadr[b]
- ndof = body_dofnum[b]
- for di in range(ndof):
- dofid = dof_start + di
- jacp, jacr = support.jac_dof(
- body_parentid,
- body_rootid,
- dof_bodyid,
- body_isdofancestor,
- subtree_com_in,
- cdof_in,
- con_pos,
- b,
- dofid,
- worldid,
+ # Interpolated flex path (trilinear)
+ if vert[side] >= 0:
+ v_adr = flex_vertadr[f] + vert[side]
+ coord = flex_vert0[v_adr]
+ cn = flex_cellnum[f]
+ cx = cn[0]
+ cy = cn[1]
+ cz = cn[2]
+ ci = wp.min(int(coord[0] * float(cx)), cx - 1)
+ ci = wp.max(ci, 0)
+ cj = wp.min(int(coord[1] * float(cy)), cy - 1)
+ cj = wp.max(cj, 0)
+ ck = wp.min(int(coord[2] * float(cz)), cz - 1)
+ ck = wp.max(ck, 0)
+ local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
+ local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
+ local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
+ local = wp.vec3(local_x, local_y, local_z)
+ ny_g = cy + 1
+ nz_g = cz + 1
+ nstart = flex_nodeadr[f]
+ for li in range(2):
+ for lj in range(2):
+ for lk in range(2):
+ w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
+ if w > 1.0e-5:
+ gi = ci + li
+ gj = cj + lj
+ gk = ck + lk
+ node_idx = gi * ny_g * nz_g + gj * nz_g + gk
+ b = body_weldid[flex_nodebodyid[nstart + node_idx]]
+ local_nb, local_bodies, local_weights = _add_weight(
+ local_nb, local_bodies, local_weights, b, w * sign, True
+ )
+ elif elem[side] >= 0:
+ e = elem[side]
+ dim_f = flex_dim[f]
+ edata_adr = flex_elemdataadr[f] + e * (dim_f + 1)
+ vert_adr_f = flex_vertadr[f]
+ total_inv_dist = float(0.0)
+ blended_coord = wp.vec3(0.0, 0.0, 0.0)
+ for vi in range(4):
+ if vi <= dim_f:
+ v_idx = flex_elem[edata_adr + vi]
+ vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx]
+ dist_v = wp.length(con_pos - vpos)
+ w_inv = 1.0 / wp.max(1.0e-10, dist_v)
+ total_inv_dist += w_inv
+ blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv
+ if total_inv_dist > 1.0e-10:
+ blended_coord = blended_coord / total_inv_dist
+ cn = flex_cellnum[f]
+ cx = cn[0]
+ cy = cn[1]
+ cz = cn[2]
+ ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1)
+ ci = wp.max(ci, 0)
+ cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1)
+ cj = wp.max(cj, 0)
+ ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1)
+ ck = wp.max(ck, 0)
+ local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0)
+ local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0)
+ local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0)
+ local = wp.vec3(local_x, local_y, local_z)
+ ny_g = cy + 1
+ nz_g = cz + 1
+ nstart = flex_nodeadr[f]
+ for li in range(2):
+ for lj in range(2):
+ for lk in range(2):
+ w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
+ if w > 1.0e-5:
+ gi = ci + li
+ gj = cj + lj
+ gk = ck + lk
+ node_idx = gi * ny_g * nz_g + gj * nz_g + gk
+ b = body_weldid[flex_nodebodyid[nstart + node_idx]]
+ local_nb, local_bodies, local_weights = _add_weight(
+ local_nb, local_bodies, local_weights, b, w * sign, True
)
-
- if wp.static(IS_ELLIPTIC):
- J = float(0.0)
- if dimid < 3:
- frame_row = frame_in[conid, dimid]
- for xyz in range(3):
- J += frame_row[xyz] * jacp[xyz] * w_sign
- else:
- frame_row = frame_in[conid, dimid - 3]
- for xyz in range(3):
- J += frame_row[xyz] * jacr[xyz] * w_sign
- else:
- J = float(0.0)
- Ji = float(0.0)
- for xyz in range(3):
- J += frame_0[xyz] * jacp[xyz] * w_sign
- if condim > 1:
- if dimid2 < 3:
- Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * w_sign
- else:
- Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * w_sign
- if condim > 1:
- if dimid % 2 == 0:
- J += Ji * frii
- else:
- J -= Ji * frii
-
- if nnz < rownnz:
- sparseid = rowadr + nnz
- efc_J_colind_out[worldid, 0, sparseid] = dofid
- efc_J_out[worldid, 0, sparseid] = J
- Jqvel += J * qvel_in[worldid, dofid]
- nnz += 1
else:
- # Non-interpolated flex: single body
- b = body_weldid[flex_vertbodyid[flex_vertadr[f] + vert[side]]]
- dof_start = body_dofadr[b]
- ndof = body_dofnum[b]
- for di in range(ndof):
- dofid = dof_start + di
- jacp, jacr = support.jac_dof(
- body_parentid,
- body_rootid,
- dof_bodyid,
- body_isdofancestor,
- subtree_com_in,
- cdof_in,
+ # Non-interpolated flex path: use pre-computed bodies/weights
+ if side == 0:
+ body_ids, weights = _get_contact_bodies_and_weights(
+ geom_bodyid,
+ flex_dim,
+ flex_cellnum,
+ flex_nodeadr,
+ flex_vertadr,
+ flex_elemdataadr,
+ flex_shelldataadr,
+ flex_nodebodyid,
+ flex_vertbodyid,
+ flex_elem,
+ flex_shell,
+ flex_vert0,
+ flexvert_xpos_in,
+ conid,
+ 0,
+ geom,
+ flex,
+ elem,
+ vert,
con_pos,
- b,
- dofid,
worldid,
)
+ else:
+ body_ids, weights = _get_contact_bodies_and_weights(
+ geom_bodyid,
+ flex_dim,
+ flex_cellnum,
+ flex_nodeadr,
+ flex_vertadr,
+ flex_elemdataadr,
+ flex_shelldataadr,
+ flex_nodebodyid,
+ flex_vertbodyid,
+ flex_elem,
+ flex_shell,
+ flex_vert0,
+ flexvert_xpos_in,
+ conid,
+ 1,
+ geom,
+ flex,
+ elem,
+ vert,
+ con_pos,
+ worldid,
+ )
+ for vi in range(4):
+ if body_ids[vi] >= 0:
+ b = body_weldid[body_ids[vi]]
+ local_nb, local_bodies, local_weights = _add_weight(
+ local_nb, local_bodies, local_weights, b, weights[vi] * sign, True
+ )
- if wp.static(IS_ELLIPTIC):
- J = float(0.0)
- if dimid < 3:
- frame_row = frame_in[conid, dimid]
- for xyz in range(3):
- J += frame_row[xyz] * jacp[xyz] * sign
+ # Evaluate Jacobians for unique bodies and write to sparse J
+ nnz = int(0)
+ Jqvel = float(0.0)
+
+ for i in range(16):
+ if i >= local_nb:
+ break
+ b = local_bodies[i]
+ w_sign = local_weights[i]
+
+ dof_start = body_dofadr[b]
+ ndof = body_dofnum[b]
+ for di in range(ndof):
+ dofid = dof_start + di
+ jacp, jacr = support.jac_dof(
+ body_parentid, body_rootid, dof_bodyid, body_isdofancestor, subtree_com_in, cdof_in, con_pos, b, dofid, worldid
+ )
+
+ if wp.static(IS_ELLIPTIC):
+ J = float(0.0)
+ if dimid < 3:
+ frame_row = frame_in[conid, dimid]
+ for xyz in range(3):
+ J += frame_row[xyz] * jacp[xyz] * w_sign
+ else:
+ frame_row = frame_in[conid, dimid - 3]
+ for xyz in range(3):
+ J += frame_row[xyz] * jacr[xyz] * w_sign
+ else:
+ J = float(0.0)
+ Ji = float(0.0)
+ for xyz in range(3):
+ J += frame_0[xyz] * jacp[xyz] * w_sign
+ if condim > 1:
+ if dimid2 < 3:
+ Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * w_sign
else:
- frame_row = frame_in[conid, dimid - 3]
- for xyz in range(3):
- J += frame_row[xyz] * jacr[xyz] * sign
+ Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * w_sign
+ if condim > 1:
+ if dimid % 2 == 0:
+ J += Ji * frii
else:
- J = float(0.0)
- Ji = float(0.0)
- for xyz in range(3):
- J += frame_0[xyz] * jacp[xyz] * sign
- if condim > 1:
- if dimid2 < 3:
- Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * sign
- else:
- Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * sign
- if condim > 1:
- if dimid % 2 == 0:
- J += Ji * frii
- else:
- J -= Ji * frii
+ J -= Ji * frii
- if nnz < rownnz:
- sparseid = rowadr + nnz
- efc_J_colind_out[worldid, 0, sparseid] = dofid
- efc_J_out[worldid, 0, sparseid] = J
- Jqvel += J * qvel_in[worldid, dofid]
- nnz += 1
-
- elif flex[side] >= 0 and elem[side] >= 0:
- # Elem contact: compute blended coordinate from distance weighting
- f = flex[side]
- e = elem[side]
- dim_f = flex_dim[f]
- edata_adr = flex_elemdataadr[f] + e * (dim_f + 1)
- vert_adr_f = flex_vertadr[f]
-
- total_inv_dist = float(0.0)
- blended_coord = wp.vec3(0.0, 0.0, 0.0)
- for vi in range(4):
- if vi <= dim_f:
- v_idx = flex_elem[edata_adr + vi]
- vpos = flexvert_xpos_in[worldid, vert_adr_f + v_idx]
- dist_v = wp.length(con_pos - vpos)
- w_inv = 1.0 / wp.max(1.0e-10, dist_v)
- total_inv_dist += w_inv
- blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv
- if total_inv_dist > 1.0e-10:
- blended_coord = blended_coord / total_inv_dist
-
- if flex_interp[f] != 0:
- cn = flex_cellnum[f]
- cx = cn[0]
- cy = cn[1]
- cz = cn[2]
- ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1)
- ci = wp.max(ci, 0)
- cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1)
- cj = wp.max(cj, 0)
- ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1)
- ck = wp.max(ck, 0)
- local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0)
- local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0)
- local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0)
- local = wp.vec3(local_x, local_y, local_z)
- ny_g = cy + 1
- nz_g = cz + 1
- nstart = flex_nodeadr[f]
-
- for li in range(2):
- for lj in range(2):
- for lk in range(2):
- w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
- if w > 1.0e-5:
- gi = ci + li
- gj = cj + lj
- gk = ck + lk
- node_idx = gi * ny_g * nz_g + gj * nz_g + gk
- b = body_weldid[flex_nodebodyid[nstart + node_idx]]
- w_sign = w * sign
-
- dof_start = body_dofadr[b]
- ndof = body_dofnum[b]
- for di in range(ndof):
- dofid = dof_start + di
- jacp, jacr = support.jac_dof(
- body_parentid,
- body_rootid,
- dof_bodyid,
- body_isdofancestor,
- subtree_com_in,
- cdof_in,
- con_pos,
- b,
- dofid,
- worldid,
- )
-
- if wp.static(IS_ELLIPTIC):
- J = float(0.0)
- if dimid < 3:
- frame_row = frame_in[conid, dimid]
- for xyz in range(3):
- J += frame_row[xyz] * jacp[xyz] * w_sign
- else:
- frame_row = frame_in[conid, dimid - 3]
- for xyz in range(3):
- J += frame_row[xyz] * jacr[xyz] * w_sign
- else:
- J = float(0.0)
- Ji = float(0.0)
- for xyz in range(3):
- J += frame_0[xyz] * jacp[xyz] * w_sign
- if condim > 1:
- if dimid2 < 3:
- Ji += frame_in[conid, dimid2][xyz] * jacp[xyz] * w_sign
- else:
- Ji += frame_in[conid, dimid2 - 3][xyz] * jacr[xyz] * w_sign
- if condim > 1:
- if dimid % 2 == 0:
- J += Ji * frii
- else:
- J -= Ji * frii
-
- if nnz < rownnz:
- sparseid = rowadr + nnz
- efc_J_colind_out[worldid, 0, sparseid] = dofid
- efc_J_out[worldid, 0, sparseid] = J
- Jqvel += J * qvel_in[worldid, dofid]
- nnz += 1
+ if nnz < rownnz:
+ sparseid = rowadr + nnz
+ efc_J_colind_out[worldid, 0, sparseid] = dofid
+ efc_J_out[worldid, 0, sparseid] = J
+ Jqvel += J * qvel_in[worldid, dofid]
+ nnz += 1
efc_Jqvel_out[worldid, efcid] = Jqvel
@@ -3752,7 +3744,7 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
TILE_SIZE = tile_size
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
body_rootid: wp.array[int],
@@ -3881,7 +3873,7 @@ def _efc_contact_jac_dense_flex(tile_size: int, cone_type: types.ConeType):
TILE_SIZE = tile_size
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
body_rootid: wp.array[int],
@@ -4197,7 +4189,7 @@ def _efc_contact_jac_dense_flex(tile_size: int, cone_type: types.ConeType):
def _efc_contact_update(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
def kernel(
# Model:
opt_timestep: wp.array[float],
@@ -4336,7 +4328,7 @@ def _efc_contact_update(cone_type: types.ConeType):
def _efc_contact_update_flex(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
opt_timestep: wp.array[float],
@@ -4432,6 +4424,11 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
elif flex[0] >= 0:
f1 = flex[0]
if flex_interp[f1] != 0:
+ # Interpolated path: de-duplicate nodes in the cell
+ local_bodies = types.vec16i(-1)
+ local_weights = types.vec16(0.0)
+ local_nb = int(0)
+
if vert[0] >= 0:
v_adr = flex_vertadr[f1] + vert[0]
coord = flex_vert0[v_adr]
@@ -4439,23 +4436,19 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
cx = cn[0]
cy = cn[1]
cz = cn[2]
-
ci = wp.min(int(coord[0] * float(cx)), cx - 1)
ci = wp.max(ci, 0)
cj = wp.min(int(coord[1] * float(cy)), cy - 1)
cj = wp.max(cj, 0)
ck = wp.min(int(coord[2] * float(cz)), cz - 1)
ck = wp.max(ck, 0)
-
local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
local = wp.vec3(local_x, local_y, local_z)
-
ny_g = cy + 1
nz_g = cz + 1
nstart = flex_nodeadr[f1]
-
for li in range(2):
for lj in range(2):
for lk in range(2):
@@ -4466,13 +4459,12 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
gk = ck + lk
node_idx = gi * ny_g * nz_g + gj * nz_g + gk
b = flex_nodebodyid[nstart + node_idx]
- invweight1 += body_invweight0[body_invweight0_id, b][0] * w
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, w, False)
elif elem[0] >= 0:
e = elem[0]
dim_f = flex_dim[f1]
edata_adr = flex_elemdataadr[f1] + e * (dim_f + 1)
vert_adr_f = flex_vertadr[f1]
-
total_inv_dist = float(0.0)
blended_coord = wp.vec3(0.0, 0.0, 0.0)
for vi in range(4):
@@ -4483,31 +4475,25 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
w_inv = 1.0 / wp.max(1.0e-10, dist_v)
total_inv_dist += w_inv
blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv
-
if total_inv_dist > 1.0e-10:
blended_coord = blended_coord / total_inv_dist
-
cn = flex_cellnum[f1]
cx = cn[0]
cy = cn[1]
cz = cn[2]
-
ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1)
ci = wp.max(ci, 0)
cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1)
cj = wp.max(cj, 0)
ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1)
ck = wp.max(ck, 0)
-
local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0)
local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0)
local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0)
local = wp.vec3(local_x, local_y, local_z)
-
ny_g = cy + 1
nz_g = cz + 1
nstart = flex_nodeadr[f1]
-
for li in range(2):
for lj in range(2):
for lk in range(2):
@@ -4518,7 +4504,14 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
gk = ck + lk
node_idx = gi * ny_g * nz_g + gj * nz_g + gk
b = flex_nodebodyid[nstart + node_idx]
- invweight1 += body_invweight0[body_invweight0_id, b][0] * w
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, w, False)
+
+ # Accumulate linear weights for unique bodies
+ for i in range(16):
+ if i < local_nb:
+ b = local_bodies[i]
+ w = local_weights[i]
+ invweight1 += body_invweight0[body_invweight0_id, b][0] * w
else:
body_ids, weights = _get_contact_bodies_and_weights(
geom_bodyid,
@@ -4563,6 +4556,11 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
elif flex[1] >= 0:
f2 = flex[1]
if flex_interp[f2] != 0:
+ # Interpolated path: de-duplicate nodes in the cell
+ local_bodies = types.vec16i(-1)
+ local_weights = types.vec16(0.0)
+ local_nb = int(0)
+
if vert[1] >= 0:
v_adr = flex_vertadr[f2] + vert[1]
coord = flex_vert0[v_adr]
@@ -4570,23 +4568,19 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
cx = cn[0]
cy = cn[1]
cz = cn[2]
-
ci = wp.min(int(coord[0] * float(cx)), cx - 1)
ci = wp.max(ci, 0)
cj = wp.min(int(coord[1] * float(cy)), cy - 1)
cj = wp.max(cj, 0)
ck = wp.min(int(coord[2] * float(cz)), cz - 1)
ck = wp.max(ck, 0)
-
local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
local = wp.vec3(local_x, local_y, local_z)
-
ny_g = cy + 1
nz_g = cz + 1
nstart = flex_nodeadr[f2]
-
for li in range(2):
for lj in range(2):
for lk in range(2):
@@ -4597,13 +4591,12 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
gk = ck + lk
node_idx = gi * ny_g * nz_g + gj * nz_g + gk
b = flex_nodebodyid[nstart + node_idx]
- invweight2 += body_invweight0[body_invweight0_id, b][0] * w
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, w, False)
elif elem[1] >= 0:
e = elem[1]
dim_f = flex_dim[f2]
edata_adr = flex_elemdataadr[f2] + e * (dim_f + 1)
vert_adr_f = flex_vertadr[f2]
-
total_inv_dist = float(0.0)
blended_coord = wp.vec3(0.0, 0.0, 0.0)
for vi in range(4):
@@ -4614,31 +4607,25 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
w_inv = 1.0 / wp.max(1.0e-10, dist_v)
total_inv_dist += w_inv
blended_coord += flex_vert0[vert_adr_f + v_idx] * w_inv
-
if total_inv_dist > 1.0e-10:
blended_coord = blended_coord / total_inv_dist
-
cn = flex_cellnum[f2]
cx = cn[0]
cy = cn[1]
cz = cn[2]
-
ci = wp.min(int(blended_coord[0] * float(cx)), cx - 1)
ci = wp.max(ci, 0)
cj = wp.min(int(blended_coord[1] * float(cy)), cy - 1)
cj = wp.max(cj, 0)
ck = wp.min(int(blended_coord[2] * float(cz)), cz - 1)
ck = wp.max(ck, 0)
-
local_x = wp.clamp(blended_coord[0] * float(cx) - float(ci), 0.0, 1.0)
local_y = wp.clamp(blended_coord[1] * float(cy) - float(cj), 0.0, 1.0)
local_z = wp.clamp(blended_coord[2] * float(cz) - float(ck), 0.0, 1.0)
local = wp.vec3(local_x, local_y, local_z)
-
ny_g = cy + 1
nz_g = cz + 1
nstart = flex_nodeadr[f2]
-
for li in range(2):
for lj in range(2):
for lk in range(2):
@@ -4649,7 +4636,14 @@ def _efc_contact_update_flex(cone_type: types.ConeType):
gk = ck + lk
node_idx = gi * ny_g * nz_g + gj * nz_g + gk
b = flex_nodebodyid[nstart + node_idx]
- invweight2 += body_invweight0[body_invweight0_id, b][0] * w
+ local_nb, local_bodies, local_weights = _add_weight(local_nb, local_bodies, local_weights, b, w, False)
+
+ # Accumulate linear weights for unique bodies
+ for i in range(16):
+ if i < local_nb:
+ b = local_bodies[i]
+ w = local_weights[i]
+ invweight2 += body_invweight0[body_invweight0_id, b][0] * w
else:
body_ids, weights = _get_contact_bodies_and_weights(
geom_bodyid,
@@ -5073,6 +5067,7 @@ def make_constraint(m: types.Model, d: types.Data):
d.njmax,
d.njmax_nnz,
d.flexnode_xpos,
+ d.face_quat,
],
outputs=[
d.ne,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
index 9329302b..681ce7a6 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
@@ -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,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
index 633c2a3a..7d531176 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
@@ -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:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
index 42149578..5bc3d2d7 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
index c6809d0e..46f0194c 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
@@ -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)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
index 090a3fc2..580d8c6a 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
@@ -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(
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
index 6c863355..4d88b2f5 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
@@ -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(
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
index d5052f9b..194a56db 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
@@ -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)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
index d9b20861..c1e697aa 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
index 1e294952..5535ca92 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
@@ -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]:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
index b87c8b2e..9fc0969c 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
index dfcc8f95..f451317b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
@@ -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],
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
index 3e9521cf..b4660590 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
@@ -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:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
index 1defb28a..1b60b95d 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
index d66ca3c9..d2dc9cfb 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
@@ -15,6 +15,7 @@
import dataclasses
from math import ceil
+from typing import Any
import warp as wp
@@ -23,14 +24,16 @@ from mujoco.mjx.third_party.mujoco_warp._src import math
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.block_cholesky import create_blocked_cholesky_augmented_factorize_solve_newton_func
from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_factorize_solve_func
-from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
+from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_newton_func
+from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import solve_search_sums
from mujoco.mjx.third_party.mujoco_warp._src.types import InverseContext
from mujoco.mjx.third_party.mujoco_warp._src.types import SolverContext
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})
_BLOCK_CHOLESKY_DIM = 32
@@ -52,8 +55,10 @@ def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext:
Jaref=wp.empty((nworld, njmax), dtype=float),
search_dot=wp.empty((nworld,), dtype=float),
done=wp.empty((nworld,), dtype=bool),
- changed_efc_ids=wp.empty((nworld, 0), dtype=int),
- changed_efc_count=wp.empty((0,), dtype=int),
+ quad_changed_ids=wp.empty((nworld, 0), dtype=int),
+ quad_changed_count=wp.empty((0,), dtype=int),
+ state_changed_count=wp.empty((0,), dtype=int),
+ ls_exhausted=wp.empty((0,), dtype=bool),
)
@@ -74,6 +79,8 @@ def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
alloc_h = m.opt.solver == types.SolverType.NEWTON
alloc_hfactor = alloc_h and nv > _BLOCK_CHOLESKY_DIM
+ alloc_mgrad = m.opt.solver == types.SolverType.CG
+ alloc_incremental = _use_incremental(m)
return SolverContext(
Jaref=wp.empty((nworld, njmax), dtype=float),
@@ -81,7 +88,8 @@ def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
done=wp.empty((nworld,), dtype=bool),
grad=wp.zeros((nworld, nv_pad), dtype=float),
grad_dot=wp.empty((nworld,), dtype=float),
- Mgrad=wp.empty((nworld, nv_pad), dtype=float),
+ newton_decrement=wp.empty((nworld,), dtype=float),
+ Mgrad=wp.empty((nworld, nv_pad), dtype=float) if alloc_mgrad else wp.empty((nworld, 0), dtype=float),
search=wp.empty((nworld, nv), dtype=float),
mv=wp.empty((nworld, nv), dtype=float),
jv=wp.empty((nworld, njmax), dtype=float),
@@ -89,14 +97,17 @@ def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
alpha=wp.empty((nworld,), dtype=float),
grad_scale=wp.empty((nworld,), dtype=float),
improvement=wp.empty((nworld,), dtype=float),
- prev_grad=wp.empty((nworld, nv), dtype=float),
- prev_Mgrad=wp.empty((nworld, nv), dtype=float),
- beta=wp.empty((nworld,), dtype=float),
- beta_den=wp.empty((nworld,), dtype=float),
+ ls_exhausted=wp.zeros((nworld,), dtype=bool),
+ search_unchanged=wp.empty((nworld,), dtype=bool),
+ prev_grad=wp.empty((nworld, nv), dtype=float) if alloc_mgrad else wp.empty((nworld, 0), dtype=float),
+ prev_Mgrad=wp.empty((nworld, nv), dtype=float) if alloc_mgrad else wp.empty((nworld, 0), dtype=float),
+ beta=wp.empty((nworld,), dtype=float) if alloc_mgrad else wp.empty((0,), dtype=float),
+ beta_den=wp.empty((nworld,), dtype=float) if alloc_mgrad else wp.empty((0,), dtype=float),
h=wp.empty((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float),
hfactor=wp.empty((nworld, nv_pad, nv_pad), dtype=float) if alloc_hfactor else wp.empty((nworld, 0, 0), dtype=float),
- changed_efc_ids=wp.empty((nworld, njmax), dtype=int) if alloc_h else wp.empty((nworld, 0), dtype=int),
- changed_efc_count=wp.empty((nworld,), dtype=int) if alloc_h else wp.empty((0,), dtype=int),
+ quad_changed_ids=wp.empty((nworld, njmax), dtype=int) if alloc_incremental else wp.empty((nworld, 0), dtype=int),
+ quad_changed_count=wp.empty((nworld,), dtype=int) if alloc_incremental else wp.empty((0,), dtype=int),
+ state_changed_count=wp.empty((nworld,), dtype=int) if alloc_incremental else wp.empty((0,), dtype=int),
)
@@ -260,13 +271,88 @@ def _eval_frictionloss_pt_3alphas(
@wp.func
-def _eval_elliptic(
+def _eval_elliptic_reference(
+ # In:
+ mu: float,
+ quad: wp.vec3,
+ quad1: wp.vec3,
+ quad2: wp.vec3,
+) -> tuple[float, float, float, int]:
+ u0 = quad1[0]
+ uu = quad1[2]
+ dm = quad2[2]
+
+ if uu <= 0.0:
+ if u0 < 0.0:
+ return quad[0], 0.0, 0.0, int(types.ConstraintState.QUADRATIC)
+ return 0.0, 0.0, 0.0, int(types.ConstraintState.SATISFIED)
+
+ T0 = wp.sqrt(uu)
+ if u0 >= mu * T0:
+ return 0.0, T0, 0.0, int(types.ConstraintState.SATISFIED)
+ if mu * u0 + T0 <= 0.0:
+ return quad[0], T0, 0.0, int(types.ConstraintState.QUADRATIC)
+
+ r0 = u0 - mu * T0
+ return 0.5 * dm * r0 * r0, T0, r0, int(types.ConstraintState.CONE)
+
+
+@wp.func
+def _eval_elliptic_alpha_zero(mu: float, quad: wp.vec3, quad1: wp.vec3, quad2: wp.vec3) -> wp.vec3:
+ cost0, T0, r0, state0 = _eval_elliptic_reference(mu, quad, quad1, quad2)
+ if state0 == int(types.ConstraintState.QUADRATIC):
+ return _eval_pt(quad, 0.0)
+ if state0 == int(types.ConstraintState.CONE):
+ T0_inv = 1.0 / T0
+ T1 = quad2[0] * T0_inv
+ T2 = (quad2[1] - T1 * T1) * T0_inv
+ r1 = quad1[1] - mu * T1
+ dm = quad2[2]
+ return wp.vec3(cost0, dm * r0 * r1, dm * (r1 * r1 - mu * r0 * T2))
+ return wp.vec3(0.0)
+
+
+@wp.func
+def _eval_elliptic_quadratic_cone_gap(mu: float, N: float, T: float, dm: float) -> float:
+ """Return quadratic cost minus cone cost at the same point."""
+ boundary = mu * N + T
+ return 0.5 * dm * boundary * boundary
+
+
+@wp.func
+def _eval_elliptic_quadratic_shifted(
+ # In:
+ mu: float,
+ quad: wp.vec3,
+ alpha: float,
+ N: float,
+ Tsqr: float,
+ u0: float,
+ T0: float,
+ dm: float,
+ state0: int,
+) -> wp.vec3:
+ aq2 = alpha * quad[2]
+ cost = alpha * (aq2 + quad[1])
+ if state0 == int(types.ConstraintState.CONE):
+ cost += _eval_elliptic_quadratic_cone_gap(mu, u0, T0, dm)
+ elif state0 == int(types.ConstraintState.SATISFIED):
+ cost = 0.5 * dm * (1.0 + mu * mu) * (N * N + wp.max(Tsqr, 0.0))
+ return wp.vec3(cost, 2.0 * aq2 + quad[1], 2.0 * quad[2])
+
+
+@wp.func
+def _eval_elliptic_shifted(
# In:
mu: float,
quad: wp.vec3,
quad1: wp.vec3,
quad2: wp.vec3,
alpha: float,
+ cost0: float,
+ T0: float,
+ r0: float,
+ state0: int,
) -> wp.vec3:
u0 = quad1[0]
v0 = quad1[1]
@@ -275,81 +361,44 @@ def _eval_elliptic(
vv = quad2[1]
dm = quad2[2]
- # compute N, Tsqr
N = u0 + alpha * v0
- Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
+ Tsqr_delta = alpha * (2.0 * uv + alpha * vv)
+ Tsqr = uu + Tsqr_delta
- # no tangential force: top or bottom zone
if Tsqr <= 0.0:
- # bottom zone: quadratic cost
if N < 0.0:
- return _eval_pt(quad, alpha)
-
- # top zone: nothing to do
- # otherwise regular processing
+ return _eval_elliptic_quadratic_shifted(mu, quad, alpha, N, Tsqr, u0, T0, dm, state0)
else:
- # tangential force
T = wp.sqrt(Tsqr)
-
- # N >= mu * T : top zone
if N >= mu * T:
- # nothing to do
pass
- # mu * N + T <= 0 : bottom zone
elif mu * N + T <= 0.0:
- return _eval_pt(quad, alpha)
-
- # otherwise middle zone
+ return _eval_elliptic_quadratic_shifted(mu, quad, alpha, N, Tsqr, u0, T0, dm, state0)
else:
- # derivatives
- N1 = v0
- T1 = (uv + alpha * vv) / T
- T2 = vv / T - (uv + alpha * vv) * T1 / (T * T)
+ T_inv = 1.0 / T
+ T1 = (uv + alpha * vv) * T_inv
+ T2 = (vv - T1 * T1) * T_inv
+ r = N - mu * T
+ r1 = v0 - mu * T1
- # add to cost
- cost = wp.vec3(
- 0.5 * dm * (N - mu * T) * (N - mu * T),
- dm * (N - mu * T) * (N1 - mu * T1),
- dm * ((N1 - mu * T1) * (N1 - mu * T1) + (N - mu * T) * (-mu * T2)),
+ if state0 == int(types.ConstraintState.CONE):
+ # Rationalize T - T0 before forming the small cone residual change.
+ T_delta = Tsqr_delta / (T + T0)
+ r_delta = alpha * v0 - mu * T_delta
+ cost = 0.5 * dm * r_delta * (2.0 * r0 + r_delta)
+ elif state0 == int(types.ConstraintState.QUADRATIC):
+ aq2 = alpha * quad[2]
+ cost = alpha * (aq2 + quad[1]) - _eval_elliptic_quadratic_cone_gap(mu, N, T, dm)
+ else:
+ cost = 0.5 * dm * r * r
+
+ return wp.vec3(
+ cost,
+ dm * r * r1,
+ dm * (r1 * r1 + r * (-mu * T2)),
)
- return cost
-
- return wp.vec3(0.0, 0.0, 0.0)
-
-
-@wp.func
-def _eval_elliptic_cost(
- # In:
- mu: float,
- quad: wp.vec3,
- quad1: wp.vec3,
- quad2: wp.vec3,
- alpha: float,
-) -> float:
- u0 = quad1[0]
- v0 = quad1[1]
- uu = quad1[2]
- uv = quad2[0]
- vv = quad2[1]
- dm = quad2[2]
-
- N = u0 + alpha * v0
- Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
-
- if Tsqr <= 0.0:
- if N < 0.0:
- return _eval_cost(quad, alpha)
- else:
- T = wp.sqrt(Tsqr)
- if N >= mu * T:
- pass
- elif mu * N + T <= 0.0:
- return _eval_cost(quad, alpha)
- else:
- return 0.5 * dm * (N - mu * T) * (N - mu * T)
-
- return 0.0
+ return wp.vec3(-cost0, 0.0, 0.0)
@wp.func
@@ -485,8 +534,8 @@ def _compute_efc_eval_pt_elliptic(
if efcid != efc_address0: # Not primary row
return wp.vec3(0.0)
mu = contact_friction[0] * impratio_invsqrt
- cost0 = _eval_elliptic_cost(mu, ctx_quad, quad1, quad2, 0.0)
- return _shift_cost(_eval_elliptic(mu, ctx_quad, quad1, quad2, alpha), cost0)
+ cost0, T0, r0, state0 = _eval_elliptic_reference(mu, ctx_quad, quad1, quad2)
+ return _eval_elliptic_shifted(mu, ctx_quad, quad1, quad2, alpha, cost0, T0, r0, state0)
# Limit/other constraint — direct eval (no quad read)
x = ctx_Jaref + alpha * ctx_jv
@@ -566,7 +615,7 @@ def _compute_efc_eval_pt_alpha_zero_elliptic(
if efcid != efc_address0: # Not primary row
return wp.vec3(0.0)
mu = contact_friction[0] * impratio_invsqrt
- return _eval_elliptic(mu, ctx_quad, quad1, quad2, 0.0)
+ return _eval_elliptic_alpha_zero(mu, ctx_quad, quad1, quad2)
# Limit/other constraint — direct eval (no quad read)
if ctx_Jaref < 0.0:
@@ -662,11 +711,6 @@ def _compute_efc_eval_pt_3alphas_elliptic(
Returns a tuple of 3 vec3s for (lo_alpha, hi_alpha, mid_alpha).
Constraint types checked in order: contact elliptic/limit/other -> friction -> equality.
"""
- # x = search point, needed for friction and limit constraints
- x_lo = ctx_Jaref + lo_alpha * ctx_jv
- x_hi = ctx_Jaref + hi_alpha * ctx_jv
- x_mid = ctx_Jaref + mid_alpha * ctx_jv
-
# Contact/limit/other constraints
if efcid >= ne + nf:
# Contact elliptic: uses special elliptic cone evaluation
@@ -674,13 +718,17 @@ def _compute_efc_eval_pt_3alphas_elliptic(
if efcid != efc_address0: # secondary rows contribute nothing
return (wp.vec3(0.0), wp.vec3(0.0), wp.vec3(0.0))
mu = contact_friction[0] * impratio_invsqrt
- cost0 = _eval_elliptic_cost(mu, ctx_quad, quad1, quad2, 0.0)
- lo = _eval_elliptic(mu, ctx_quad, quad1, quad2, lo_alpha)
- hi = _eval_elliptic(mu, ctx_quad, quad1, quad2, hi_alpha)
- mid = _eval_elliptic(mu, ctx_quad, quad1, quad2, mid_alpha)
- return (_shift_cost(lo, cost0), _shift_cost(hi, cost0), _shift_cost(mid, cost0))
+ cost0, T0, r0, state0 = _eval_elliptic_reference(mu, ctx_quad, quad1, quad2)
+ return (
+ _eval_elliptic_shifted(mu, ctx_quad, quad1, quad2, lo_alpha, cost0, T0, r0, state0),
+ _eval_elliptic_shifted(mu, ctx_quad, quad1, quad2, hi_alpha, cost0, T0, r0, state0),
+ _eval_elliptic_shifted(mu, ctx_quad, quad1, quad2, mid_alpha, cost0, T0, r0, state0),
+ )
# Limit/other constraints — direct eval (no quad read)
+ x_lo = ctx_Jaref + lo_alpha * ctx_jv
+ x_hi = ctx_Jaref + hi_alpha * ctx_jv
+ x_mid = ctx_Jaref + mid_alpha * ctx_jv
efc_D = efc_D_in[efcid]
quad0 = _eval_pt_direct_cost_alpha_zero(ctx_Jaref, efc_D)
cost0 = wp.where(ctx_Jaref < 0.0, quad0, 0.0)
@@ -696,6 +744,9 @@ def _compute_efc_eval_pt_3alphas_elliptic(
# Friction constraint - load D and frictionloss only here
if efcid >= ne:
+ x_lo = ctx_Jaref + lo_alpha * ctx_jv
+ x_hi = ctx_Jaref + hi_alpha * ctx_jv
+ x_mid = ctx_Jaref + mid_alpha * ctx_jv
efc_D = efc_D_in[efcid]
f = efc_frictionloss[efcid]
rf = math.safe_div(f, efc_D)
@@ -769,7 +820,9 @@ def _compute_efc_eval_pt_3alphas_elliptic(
@cache_kernel
-def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool, is_sparse: bool):
+def _linesearch_iterative_kernel(
+ ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool, is_sparse: bool, incremental: bool
+):
"""Factory for iterative linesearch kernel.
Args:
@@ -777,10 +830,12 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
cone_type: Friction cone type (PYRAMIDAL or ELLIPTIC) for compile-time optimization.
fuse_jv: Whether to compute jv = J @ search in-kernel (efficient for small nv).
is_sparse: Use sparse matrix representation for constraint Jacobian.
+ incremental: State changes are tracked: flag exhausted rays, reuse jv on unchanged search.
"""
LS_ITERATIONS = ls_iterations
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
FUSE_JV = fuse_jv
+ INCREMENTAL = incremental
IS_SPARSE = is_sparse
# Native snippet for CUDA __syncthreads()
@@ -798,7 +853,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
_compute_efc_eval_pt_alpha_zero = _compute_efc_eval_pt_alpha_zero_pyramidal
_compute_efc_eval_pt_3alphas = _compute_efc_eval_pt_3alphas_pyramidal
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
def kernel(
# Model:
nv: int,
@@ -825,6 +880,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
njmax_in: int,
nacon_in: wp.array[int],
# In:
+ ctx_search_unchanged_in: wp.array[bool],
ctx_Jaref_in: wp.array2d[float],
ctx_search_in: wp.array2d[float],
ctx_search_dot_in: wp.array[float],
@@ -841,6 +897,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
ctx_quad_out: wp.array2d[wp.vec3],
ctx_improvement_out: wp.array[float],
ctx_alpha_out: wp.array[float],
+ ctx_ls_exhausted_out: wp.array[bool],
):
worldid, tid = wp.tid()
@@ -851,21 +908,26 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
nf = nf_in[worldid]
nefc = wp.min(njmax_in, nefc_in[worldid])
- # jv = J @ search (fused for small nv)
+ # jv = J @ search (fused for small nv); on unchanged search the buffer
+ # already holds jv (see _linesearch).
if wp.static(FUSE_JV):
- for efcid in range(tid, nefc, wp.block_dim()):
- jv = float(0.0)
- if wp.static(IS_SPARSE):
- rownnz = efc_J_rownnz_in[worldid, efcid]
- rowadr = efc_J_rowadr_in[worldid, efcid]
- for k in range(rownnz):
- sparseid = rowadr + k
- colind = efc_J_colind_in[worldid, 0, sparseid]
- jv += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
- else:
- for i in range(nv):
- jv += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i]
- ctx_jv_out[worldid, efcid] = jv
+ recompute_jv = True
+ if wp.static(INCREMENTAL):
+ recompute_jv = not ctx_search_unchanged_in[worldid]
+ if recompute_jv:
+ for efcid in range(tid, nefc, wp.block_dim()):
+ jv = float(0.0)
+ if wp.static(IS_SPARSE):
+ rownnz = efc_J_rownnz_in[worldid, efcid]
+ rowadr = efc_J_rowadr_in[worldid, efcid]
+ for k in range(rownnz):
+ sparseid = rowadr + k
+ colind = efc_J_colind_in[worldid, 0, sparseid]
+ jv += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
+ else:
+ for i in range(nv):
+ jv += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i]
+ ctx_jv_out[worldid, efcid] = jv
_syncthreads() # ensure all jv values are written before reading
@@ -1014,6 +1076,20 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
p0 = wp.vec3(ctx_quad_gauss[0], ctx_quad_gauss[1], 2.0 * ctx_quad_gauss[2]) + p0_sum[0]
p0_delta = wp.vec3(0.0, p0[1], p0[2])
+ # The bracketing search reads derivative sums whose rounding noise is ~eps
+ # times their gross magnitude, so roots below eps * |q1| / p0[2] (and never
+ # below 8 ulps of the unit ray anchor) are noise, not descent: an accepted
+ # step under this floor means the stale ray is exhausted and the fast path
+ # must rebuild the world. Per quadratic row |q1| = sqrt(2 * cost * hessian),
+ # so Cauchy-Schwarz bounds the row total by sums already reduced for p0; the
+ # smooth term is added exactly. Friction linear rows fall outside the bound,
+ # which only lowers the floor toward the fixed 8-ulp base.
+ noise_floor = float(0.0)
+ if wp.static(INCREMENTAL):
+ rows = p0_sum[0]
+ q1_abs = wp.sqrt(2.0 * wp.max(rows[0], 0.0) * wp.max(rows[2], 0.0)) + wp.abs(ctx_quad_gauss[1])
+ noise_floor = _ALPHA_NOISE_EPS * wp.max(1.0, math.safe_div(q1_abs, p0[2]))
+
# lo_in at lo_alpha_in = -p0[1] / p0[2]
lo_alpha_in = -math.safe_div(p0[1], p0[2])
@@ -1236,6 +1312,8 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
if tid == 0:
ctx_improvement_out[worldid] = improvement
ctx_alpha_out[worldid] = alpha
+ if wp.static(INCREMENTAL):
+ ctx_ls_exhausted_out[worldid] = wp.abs(alpha) < noise_floor
return kernel
@@ -1250,7 +1328,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
fuse_jv: Whether jv is computed in-kernel (True) or pre-computed (False).
"""
wp.launch_tiled(
- _linesearch_iterative_kernel(m.opt.ls_iterations, m.opt.cone, fuse_jv, m.is_sparse),
+ _linesearch_iterative_kernel(m.opt.ls_iterations, m.opt.cone, fuse_jv, m.is_sparse, _use_incremental(m)),
dim=d.nworld,
inputs=[
m.nv,
@@ -1275,6 +1353,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
d.efc.frictionloss,
d.njmax,
d.nacon,
+ ctx.search_unchanged,
ctx.Jaref,
ctx.search,
ctx.search_dot,
@@ -1283,7 +1362,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
ctx.quad,
ctx.done,
],
- outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement, ctx.alpha],
+ outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement, ctx.alpha, ctx.ls_exhausted],
block_dim=m.block_dim.linesearch_iterative,
)
@@ -1293,7 +1372,7 @@ def _linesearch_zero_jv(
# Data in:
nefc_in: wp.array[int],
# In:
- ctx_done_in: wp.array[bool],
+ skip_in: wp.array[bool],
# Out:
ctx_jv_out: wp.array2d[float],
):
@@ -1302,15 +1381,17 @@ def _linesearch_zero_jv(
if efcid >= nefc_in[worldid]:
return
- if ctx_done_in[worldid]:
+ if skip_in[worldid]:
return
ctx_jv_out[worldid, efcid] = 0.0
@cache_kernel
-def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
- @wp.kernel(module="unique", enable_backward=False)
+def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int, compact: bool):
+ COMPACT = compact
+
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Data in:
nefc_in: wp.array[int],
@@ -1318,9 +1399,10 @@ def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
efc_J_in: wp.array3d[float],
+ dof_cdof_in: wp.array2d[int],
# In:
ctx_search_in: wp.array2d[float],
- ctx_done_in: wp.array[bool],
+ skip_in: wp.array[bool],
# Out:
ctx_jv_out: wp.array2d[float],
):
@@ -1329,7 +1411,7 @@ def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
if efcid >= nefc_in[worldid]:
return
- if ctx_done_in[worldid]:
+ if skip_in[worldid]:
return
jv_out = float(0.0)
@@ -1342,6 +1424,10 @@ def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
for k in range(rownnz):
sparseid = rowadr + k
colind = efc_J_colind_in[worldid, 0, sparseid]
+ if wp.static(COMPACT):
+ colind = dof_cdof_in[worldid, colind]
+ if colind < 0:
+ continue
jv_out += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
else:
for i in range(wp.static(min(dofs_per_thread, nv))):
@@ -1357,6 +1443,10 @@ def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
for k in range(rownnz):
sparseid = rowadr + k
colind = efc_J_colind_in[worldid, 0, sparseid]
+ if wp.static(COMPACT):
+ colind = dof_cdof_in[worldid, colind]
+ if colind < 0:
+ continue
jv_out += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
ctx_jv_out[worldid, efcid] = jv_out
else:
@@ -1373,21 +1463,34 @@ def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext):
"""Linesearch for constraint solver.
+ When state changes are tracked, worlds with ctx.search_unchanged reuse last
+ iteration's mv/jv, so a fresh search requires clearing the flag (see the
+ invalidation in _solve and the writer in _update_gradient_zero_grad_dot).
+
Args:
m: Model
d: Data
ctx: SolverContext
"""
+ # mv and jv are pure functions of the search direction, and M and J are
+ # constant within a solve, so worlds whose search was kept reuse last
+ # iteration's values.
+ skip = ctx.search_unchanged if _use_incremental(m) else ctx.done
+
# mv = M @ search (common to both parallel and iterative)
- support.mul_m(m, d, ctx.mv, ctx.search, skip=ctx.done)
+ _mul_m_compact_aware(m, d, ctx, ctx.mv, ctx.search, skip)
# Fuse jv computation in-kernel for small nv (iterative only, dense only)
# Sparse mode requires pre-computed jv since in-kernel uses dense indexing
- fuse_jv = m.nv <= 50 and not m.is_sparse
+ # the sparse-compact J path reads the full model's sparse J structures;
+ # dense full models keep the gathered dense cJ path
+ sc = _sparse_compact(ctx)
+ fuse_jv = m.nv <= 50 and not m.is_sparse and not sc
# jv = J @ search (when not fused into iterative kernel)
if not fuse_jv:
- if m.is_sparse:
+ dj = ctx.compact_d_full if sc else d
+ if sc or m.is_sparse:
# Sparse J has few nonzeros per row, one thread handles them all.
dofs_per_thread = m.nv
threads_per_efc = 1
@@ -1399,21 +1502,50 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext):
wp.launch(
_linesearch_zero_jv,
dim=(d.nworld, d.njmax),
- inputs=[d.nefc, ctx.done],
+ inputs=[d.nefc, skip],
outputs=[ctx.jv],
)
wp.launch(
- _linesearch_jv_fused_kernel(m.is_sparse, m.nv, dofs_per_thread),
+ _linesearch_jv_fused_kernel(sc or m.is_sparse, m.nv, dofs_per_thread, sc),
dim=(d.nworld, d.njmax, threads_per_efc),
- inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, ctx.search, ctx.done],
+ inputs=[d.nefc, dj.efc.J_rownnz, dj.efc.J_rowadr, dj.efc.J_colind, dj.efc.J, dj.dof_cdof, ctx.search, skip],
outputs=[ctx.jv],
)
_linesearch_iterative(m, d, ctx, fuse_jv)
-@wp.kernel
+@cache_kernel
+def _solve_init_dof(warmstart: bool, sparse: bool):
+ WARMSTART = warmstart
+ SPARSE = sparse
+
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
+ def kernel(
+ # Data in:
+ nefc_in: wp.array[int],
+ qacc_warmstart_in: wp.array2d[float],
+ qacc_smooth_in: wp.array2d[float],
+ # Data out:
+ qacc_out: wp.array2d[float],
+ qfrc_constraint_out: wp.array2d[float],
+ ):
+ worldid, dofid = wp.tid()
+
+ if wp.static(WARMSTART):
+ qacc_out[worldid, dofid] = qacc_warmstart_in[worldid, dofid]
+ else:
+ qacc_out[worldid, dofid] = qacc_smooth_in[worldid, dofid]
+
+ if wp.static(SPARSE):
+ if nefc_in[worldid] == 0:
+ qfrc_constraint_out[worldid, dofid] = 0.0
+
+ return kernel
+
+
+@wp.kernel(grid_stride=True)
def _solve_init_efc(
# Data out:
solver_niter_out: wp.array[int],
@@ -1428,8 +1560,10 @@ def _solve_init_efc(
@cache_kernel
-def _solve_init_jaref_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
- @wp.kernel(module="unique", enable_backward=False)
+def _solve_init_jaref_kernel(is_sparse: bool, nv: int, dofs_per_thread: int, compact: bool):
+ COMPACT = compact
+
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
def kernel(
# Data in:
nefc_in: wp.array[int],
@@ -1439,6 +1573,7 @@ def _solve_init_jaref_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
efc_J_colind_in: wp.array3d[int],
efc_J_in: wp.array3d[float],
efc_aref_in: wp.array2d[float],
+ dof_cdof_in: wp.array2d[int],
# Out:
ctx_Jaref_out: wp.array2d[float],
):
@@ -1454,6 +1589,10 @@ def _solve_init_jaref_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
for i in range(rownnz):
sparseid = rowadr + i
colind = efc_J_colind_in[worldid, 0, sparseid]
+ if wp.static(COMPACT):
+ colind = dof_cdof_in[worldid, colind]
+ if colind < 0:
+ continue
jaref += efc_J_in[worldid, 0, sparseid] * qacc_in[worldid, colind]
ctx_Jaref_out[worldid, efcid] = jaref - efc_aref_in[worldid, efcid]
else:
@@ -1476,20 +1615,6 @@ def _solve_init_jaref_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
return kernel
-@wp.kernel
-def _solve_init_search(
- # In:
- ctx_Mgrad_in: wp.array2d[float],
- # Out:
- ctx_search_out: wp.array2d[float],
- ctx_search_dot_out: wp.array[float],
-):
- worldid, dofid = wp.tid()
- search = -1.0 * ctx_Mgrad_in[worldid, dofid]
- ctx_search_out[worldid, dofid] = search
- wp.atomic_add(ctx_search_dot_out, worldid, search * search)
-
-
@wp.kernel
def _solve_init_search_cg_tiled(
# Model:
@@ -1528,7 +1653,7 @@ def _solve_init_search_cg_tiled(
def _update_constraint_efc(track_changes: bool):
TRACK_CHANGES = track_changes
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Model:
opt_impratio_invsqrt: wp.array[float],
@@ -1546,25 +1671,33 @@ def _update_constraint_efc(track_changes: bool):
nacon_in: wp.array[int],
# In:
ctx_Jaref_in: wp.array2d[float],
+ ctx_ls_exhausted_in: wp.array[bool],
ctx_done_in: wp.array[bool],
# Data out:
efc_force_out: wp.array2d[float],
efc_state_out: wp.array2d[int],
# Out:
- changed_ids_out: wp.array2d[int],
- changed_count_out: wp.array[int],
+ quad_changed_ids_out: wp.array2d[int],
+ quad_changed_count_out: wp.array[int],
+ state_changed_count_out: wp.array[int],
):
worldid, efcid = wp.tid()
- if efcid >= nefc_in[worldid]:
- return
-
if ctx_done_in[worldid]:
return
- # Read old QUADRATIC status before overwriting
+ # The linesearch flags worlds whose accepted step was rounding noise (stale
+ # ray exhausted); count it as a state change so the fast path rebuilds them.
if wp.static(TRACK_CHANGES):
- old_quad = efc_state_out[worldid, efcid] == types.ConstraintState.QUADRATIC.value
+ if efcid == 0 and ctx_ls_exhausted_in[worldid]:
+ wp.atomic_add(state_changed_count_out, worldid, 1)
+
+ if efcid >= nefc_in[worldid]:
+ return
+
+ # Read old state before overwriting
+ if wp.static(TRACK_CHANGES):
+ old_state = efc_state_out[worldid, efcid]
efc_D = efc_D_in[worldid, efcid]
Jaref = ctx_Jaref_in[worldid, efcid]
@@ -1630,51 +1763,87 @@ def _update_constraint_efc(track_changes: bool):
efc_state_out[worldid, efcid] = new_state
if wp.static(TRACK_CHANGES):
+ old_quad = old_state == types.ConstraintState.QUADRATIC.value
new_quad = new_state == types.ConstraintState.QUADRATIC.value
if old_quad != new_quad:
- idx = wp.atomic_add(changed_count_out, worldid, 1)
- changed_ids_out[worldid, idx] = efcid
+ idx = wp.atomic_add(quad_changed_count_out, worldid, 1)
+ quad_changed_ids_out[worldid, idx] = efcid
+ # LINEARNEG <-> LINEARPOS friction transitions change the force without
+ # changing the quadratic flag (or H); the fast path must still see them.
+ if old_state != new_state:
+ wp.atomic_add(state_changed_count_out, worldid, 1)
return kernel
@wp.kernel
-def _update_constraint_init_qfrc_constraint_sparse(
- # Data in:
- nefc_in: wp.array[int],
- efc_J_rownnz_in: wp.array2d[int],
- efc_J_rowadr_in: wp.array2d[int],
- efc_J_colind_in: wp.array3d[int],
- efc_J_in: wp.array3d[float],
- efc_force_in: wp.array2d[float],
+def _zero_qfrc_constraint_sparse(
# In:
- changed_count_in: wp.array[int],
+ state_changed_count_in: wp.array[int],
ctx_done_in: wp.array[bool],
# Data out:
qfrc_constraint_out: wp.array2d[float],
):
- worldid, efcid = wp.tid()
+ # Only zero worlds the rebuild will repopulate; done worlds keep their value.
+ worldid, dofid = wp.tid()
if ctx_done_in[worldid]:
return
- if changed_count_in[worldid] == 0:
+ if state_changed_count_in[worldid] == 0:
return
- if efcid >= nefc_in[worldid]:
- return
+ qfrc_constraint_out[worldid, dofid] = 0.0
- force = efc_force_in[worldid, efcid]
- if force == 0.0:
- return
- rownnz = efc_J_rownnz_in[worldid, efcid]
- rowadr = efc_J_rowadr_in[worldid, efcid]
- for i in range(rownnz):
- sparseid = rowadr + i
- colind = efc_J_colind_in[worldid, 0, sparseid]
- efc_J = efc_J_in[worldid, 0, sparseid]
- wp.atomic_add(qfrc_constraint_out[worldid], colind, efc_J * force)
+@cache_kernel
+def _update_constraint_init_qfrc_constraint_sparse(compact: bool):
+ COMPACT = compact
+
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
+ def kernel(
+ # Data in:
+ nefc_in: wp.array[int],
+ efc_J_rownnz_in: wp.array2d[int],
+ efc_J_rowadr_in: wp.array2d[int],
+ efc_J_colind_in: wp.array3d[int],
+ efc_J_in: wp.array3d[float],
+ efc_force_in: wp.array2d[float],
+ dof_cdof_in: wp.array2d[int],
+ # In:
+ state_changed_count_in: wp.array[int],
+ ctx_done_in: wp.array[bool],
+ # Data out:
+ qfrc_constraint_out: wp.array2d[float],
+ ):
+ worldid, efcid = wp.tid()
+
+ if ctx_done_in[worldid]:
+ return
+
+ if state_changed_count_in[worldid] == 0:
+ return
+
+ if efcid >= nefc_in[worldid]:
+ return
+
+ force = efc_force_in[worldid, efcid]
+ if force == 0.0:
+ return
+
+ rownnz = efc_J_rownnz_in[worldid, efcid]
+ rowadr = efc_J_rowadr_in[worldid, efcid]
+ for i in range(rownnz):
+ sparseid = rowadr + i
+ colind = efc_J_colind_in[worldid, 0, sparseid]
+ if wp.static(COMPACT):
+ colind = dof_cdof_in[worldid, colind]
+ if colind < 0:
+ continue
+ efc_J = efc_J_in[worldid, 0, sparseid]
+ wp.atomic_add(qfrc_constraint_out[worldid], colind, efc_J * force)
+
+ return kernel
@wp.kernel
@@ -1698,7 +1867,7 @@ def _qfrc_constraint_from_grad(
def _update_constraint_init_qfrc_constraint_dense(stable_fast: bool):
STABLE_FAST = stable_fast
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Data in:
nefc_in: wp.array[int],
@@ -1706,7 +1875,7 @@ def _update_constraint_init_qfrc_constraint_dense(stable_fast: bool):
efc_force_in: wp.array2d[float],
njmax_in: int,
# In:
- changed_count_in: wp.array[int],
+ state_changed_count_in: wp.array[int],
ctx_done_in: wp.array[bool],
# Data out:
qfrc_constraint_out: wp.array2d[float],
@@ -1718,7 +1887,7 @@ def _update_constraint_init_qfrc_constraint_dense(stable_fast: bool):
# Fast path: stale qfrc_constraint is never read; recovered after the solve.
if wp.static(STABLE_FAST):
- if changed_count_in[worldid] == 0:
+ if state_changed_count_in[worldid] == 0:
return
sum_qfrc = float(0.0)
@@ -1739,8 +1908,8 @@ def _update_gradient_h_incremental(
efc_D_in: wp.array2d[float],
efc_state_in: wp.array2d[int],
# In:
- changed_ids_in: wp.array2d[int],
- changed_count_in: wp.array[int],
+ quad_changed_ids_in: wp.array2d[int],
+ quad_changed_count_in: wp.array[int],
# Out:
ctx_h_out: wp.array3d[float],
):
@@ -1751,7 +1920,7 @@ def _update_gradient_h_incremental(
"""
worldid, elementid = wp.tid()
- n_changes = changed_count_in[worldid]
+ n_changes = quad_changed_count_in[worldid]
if n_changes == 0:
return
@@ -1761,7 +1930,7 @@ def _update_gradient_h_incremental(
delta = float(0.0)
for change_idx in range(n_changes):
- efcid = changed_ids_in[worldid, change_idx]
+ efcid = quad_changed_ids_in[worldid, change_idx]
Jrow = efc_J_in[worldid, efcid, row]
if Jrow == 0.0:
continue
@@ -1779,54 +1948,66 @@ def _update_gradient_h_incremental(
ctx_h_out[worldid, row, col] += delta
-@wp.kernel
-def _update_gradient_h_incremental_sparse(
- # Data in:
- efc_J_rownnz_in: wp.array2d[int],
- efc_J_rowadr_in: wp.array2d[int],
- efc_J_colind_in: wp.array3d[int],
- efc_J_in: wp.array3d[float],
- efc_D_in: wp.array2d[float],
- efc_state_in: wp.array2d[int],
- # In:
- changed_ids_in: wp.array2d[int],
- changed_count_in: wp.array[int],
- slots_per_world: int,
- # Out:
- ctx_h_out: wp.array3d[float],
-):
- """Incrementally update upper triangle of H for changed constraints (sparse J).
+@cache_kernel
+def _update_gradient_h_incremental_sparse(compact: bool):
+ COMPACT = compact
- One warp per changed constraint row: the lanes split the row's upper-triangular
- entries (same sqrt triangular-number decode as _JTDAJ_sparse), replacing the
- serial nnz^2 loop that dominated this kernel.
- """
- worldid, slot, lane = wp.tid()
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
+ def kernel(
+ # Data in:
+ efc_J_rownnz_in: wp.array2d[int],
+ efc_J_rowadr_in: wp.array2d[int],
+ efc_J_colind_in: wp.array3d[int],
+ efc_J_in: wp.array3d[float],
+ efc_D_in: wp.array2d[float],
+ efc_state_in: wp.array2d[int],
+ dof_cdof_in: wp.array2d[int],
+ # In:
+ quad_changed_ids_in: wp.array2d[int],
+ quad_changed_count_in: wp.array[int],
+ slots_per_world: int,
+ # Out:
+ ctx_h_out: wp.array3d[float],
+ ):
+ """Incrementally update upper triangle of H for changed constraints (sparse J).
- n_changes = changed_count_in[worldid]
- for change_idx in range(slot, n_changes, slots_per_world):
- efcid = changed_ids_in[worldid, change_idx]
- D = efc_D_in[worldid, efcid]
- sign = float(0.0)
- if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value:
- sign = D
- else:
- sign = -D
+ One warp per changed constraint row: the lanes split the row's upper-triangular
+ entries (same sqrt triangular-number decode as _JTDACJ_sparse), replacing the
+ serial nnz^2 loop that dominated this kernel.
+ """
+ worldid, slot, lane = wp.tid()
- rownnz = efc_J_rownnz_in[worldid, efcid]
- rowadr = efc_J_rowadr_in[worldid, efcid]
- n_entries = rownnz * (rownnz + 1) // 2
+ n_changes = quad_changed_count_in[worldid]
+ for change_idx in range(slot, n_changes, slots_per_world):
+ efcid = quad_changed_ids_in[worldid, change_idx]
+ D = efc_D_in[worldid, efcid]
+ sign = float(0.0)
+ if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value:
+ sign = D
+ else:
+ sign = -D
- for entry in range(lane, n_entries, wp.static(_JTDAJ_THREADS_PER_GROUP)):
- ii = int((wp.sqrt(float(8 * entry + 1)) - 1.0) * 0.5)
- jj = entry - ii * (ii + 1) // 2
- Ji = efc_J_in[worldid, 0, rowadr + ii]
- Jj = efc_J_in[worldid, 0, rowadr + jj]
- h = sign * Ji * Jj
- if h != 0.0:
- colindi = efc_J_colind_in[worldid, 0, rowadr + ii]
- colindj = efc_J_colind_in[worldid, 0, rowadr + jj]
- wp.atomic_add(ctx_h_out[worldid, wp.min(colindi, colindj)], wp.max(colindi, colindj), h)
+ rownnz = efc_J_rownnz_in[worldid, efcid]
+ rowadr = efc_J_rowadr_in[worldid, efcid]
+ n_entries = rownnz * (rownnz + 1) // 2
+
+ for entry in range(lane, n_entries, wp.static(_JTDAJ_THREADS_PER_GROUP)):
+ ii = int((wp.sqrt(float(8 * entry + 1)) - 1.0) * 0.5)
+ jj = entry - ii * (ii + 1) // 2
+ Ji = efc_J_in[worldid, 0, rowadr + ii]
+ Jj = efc_J_in[worldid, 0, rowadr + jj]
+ h = sign * Ji * Jj
+ if h != 0.0:
+ colindi = efc_J_colind_in[worldid, 0, rowadr + ii]
+ colindj = efc_J_colind_in[worldid, 0, rowadr + jj]
+ if wp.static(COMPACT):
+ colindi = dof_cdof_in[worldid, colindi]
+ colindj = dof_cdof_in[worldid, colindj]
+ if colindi < 0 or colindj < 0:
+ continue
+ wp.atomic_add(ctx_h_out[worldid, wp.min(colindi, colindj)], wp.max(colindi, colindj), h)
+
+ return kernel
def _update_constraint(
@@ -1851,6 +2032,7 @@ def _update_constraint(
d.efc.frictionloss,
d.nacon,
ctx.Jaref,
+ ctx.ls_exhausted,
ctx.done,
]
@@ -1858,18 +2040,25 @@ def _update_constraint(
_update_constraint_efc(track_changes),
dim=(d.nworld, d.njmax),
inputs=efc_inputs,
- outputs=[d.efc.force, d.efc.state, ctx.changed_efc_ids, ctx.changed_efc_count],
+ outputs=[d.efc.force, d.efc.state, ctx.quad_changed_ids, ctx.quad_changed_count, ctx.state_changed_count],
)
# qfrc_constraint = efc_J.T @ efc_force. Fast-path worlds with no state flips
# skip the rebuild; the public value is recovered after the solve.
- changed = ctx.changed_efc_count if stable_fast else d.nefc
- if m.is_sparse:
- d.qfrc_constraint.zero_()
+ changed = ctx.state_changed_count if stable_fast else d.nefc
+ sc = _sparse_compact(ctx)
+ if m.is_sparse or sc:
+ dj = ctx.compact_d_full if sc else d
wp.launch(
- _update_constraint_init_qfrc_constraint_sparse,
+ _zero_qfrc_constraint_sparse,
+ dim=(d.nworld, m.nv),
+ inputs=[changed, ctx.done],
+ outputs=[d.qfrc_constraint],
+ )
+ wp.launch(
+ _update_constraint_init_qfrc_constraint_sparse(sc),
dim=(d.nworld, d.njmax),
- inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, changed, ctx.done],
+ inputs=[d.nefc, dj.efc.J_rownnz, dj.efc.J_rowadr, dj.efc.J_colind, dj.efc.J, d.efc.force, dj.dof_cdof, changed, ctx.done],
outputs=[d.qfrc_constraint],
)
else:
@@ -1885,18 +2074,25 @@ def _update_constraint(
def _update_gradient_zero_grad_dot(stable_fast: bool):
STABLE_FAST = stable_fast
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# In:
- changed_count_in: wp.array[int],
+ state_changed_count_in: wp.array[int],
ctx_alpha_in: wp.array[float],
ctx_done_in: wp.array[bool],
# Out:
ctx_grad_dot_out: wp.array[float],
+ ctx_newton_decrement_out: wp.array[float],
ctx_grad_scale_out: wp.array[float],
+ ctx_search_unchanged_out: wp.array[bool],
):
worldid = wp.tid()
+ if wp.static(STABLE_FAST):
+ ctx_search_unchanged_out[worldid] = ctx_done_in[worldid] or state_changed_count_in[worldid] == 0
+ else:
+ ctx_search_unchanged_out[worldid] = False
+
if ctx_done_in[worldid]:
return
@@ -1904,17 +2100,20 @@ def _update_gradient_zero_grad_dot(stable_fast: bool):
# gradient is grad_scale * g, and a linesearch step t along the (equally
# stale) search direction changes it to (grad_scale - t) * g.
if wp.static(STABLE_FAST):
- if changed_count_in[worldid] == 0:
+ if state_changed_count_in[worldid] == 0:
sigma = ctx_grad_scale_out[worldid]
new_sigma = sigma - ctx_alpha_in[worldid]
ratio = float(0.0)
if sigma != 0.0:
ratio = new_sigma / sigma
- ctx_grad_dot_out[worldid] *= ratio * ratio
+ ratio_sq = ratio * ratio
+ ctx_grad_dot_out[worldid] *= ratio_sq
+ ctx_newton_decrement_out[worldid] *= ratio_sq
ctx_grad_scale_out[worldid] = new_sigma
return
ctx_grad_dot_out[worldid] = 0.0
+ ctx_newton_decrement_out[worldid] = 0.0
ctx_grad_scale_out[worldid] = 1.0
return kernel
@@ -1924,14 +2123,14 @@ def _update_gradient_zero_grad_dot(stable_fast: bool):
def _update_gradient_grad(stable_fast: bool):
STABLE_FAST = stable_fast
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# Data in:
qfrc_smooth_in: wp.array2d[float],
qfrc_constraint_in: wp.array2d[float],
efc_Ma_in: wp.array2d[float],
# In:
- changed_count_in: wp.array[int],
+ state_changed_count_in: wp.array[int],
ctx_done_in: wp.array[bool],
# Out:
ctx_grad_out: wp.array2d[float],
@@ -1944,7 +2143,7 @@ def _update_gradient_grad(stable_fast: bool):
# Fast path: grad stays stale (see _update_gradient_zero_grad_dot).
if wp.static(STABLE_FAST):
- if changed_count_in[worldid] == 0:
+ if state_changed_count_in[worldid] == 0:
return
grad = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - qfrc_constraint_in[worldid, dofid]
@@ -1988,37 +2187,54 @@ def _update_gradient_grad_tiled(
ctx_grad_dot_out[worldid] = grad_dot_sum[0]
-@wp.kernel
-def _update_gradient_init_h_sparse(
- # Model:
- nv: int,
- M_elemid: wp.array2d[int],
- # Data in:
- M_in: wp.array2d[float],
- # In:
- ctx_done_in: wp.array[bool],
- # Out:
- ctx_h_out: wp.array3d[float],
-):
- worldid, i, j = wp.tid()
+@cache_kernel
+def _update_gradient_init_h_sparse(compact: bool):
+ COMPACT = compact
- if ctx_done_in[worldid]:
- return
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
+ def kernel(
+ # Model:
+ nv: int,
+ M_elemid: wp.array2d[int],
+ # Data in:
+ M_in: wp.array2d[float],
+ cdof_dof_in: wp.array2d[int],
+ # In:
+ ctx_done_in: wp.array[bool],
+ # Out:
+ ctx_h_out: wp.array3d[float],
+ ):
+ worldid, i, j = wp.tid()
- # only write the upper triangle; Cholesky reads the upper triangle only
- if j < i:
- return
+ if ctx_done_in[worldid]:
+ return
- if i >= nv or j >= nv:
- ctx_h_out[worldid, i, j] = 0.0
- return
+ # only write the upper triangle; Cholesky reads the upper triangle only
+ if j < i:
+ return
- # sparse M is stored in the lower triangle, so transpose the lookup for the upper
- elemid = M_elemid[j, i]
- if elemid >= 0:
- ctx_h_out[worldid, i, j] = M_in[worldid, elemid]
- else:
- ctx_h_out[worldid, i, j] = 0.0
+ if wp.static(COMPACT):
+ dof_i = cdof_dof_in[worldid, i]
+ dof_j = cdof_dof_in[worldid, j]
+ if dof_i < 0 or dof_j < 0:
+ # per-world padded block: identity keeps the factorization well conditioned
+ ctx_h_out[worldid, i, j] = wp.where(i == j, 1.0, 0.0)
+ return
+ else:
+ dof_i = i
+ dof_j = j
+ if i >= nv or j >= nv:
+ ctx_h_out[worldid, i, j] = 0.0
+ return
+
+ # sparse M is stored in the lower triangle, so look up (larger, smaller)
+ elemid = M_elemid[wp.max(dof_i, dof_j), wp.min(dof_i, dof_j)]
+ if elemid >= 0:
+ ctx_h_out[worldid, i, j] = M_in[worldid, elemid]
+ else:
+ ctx_h_out[worldid, i, j] = 0.0
+
+ return kernel
@wp.func
@@ -2051,7 +2267,7 @@ def _update_gradient_JTDAJ_dense_tiled_compact(nv_pad: int, tile_size: int, njma
TILE_SIZE_K = tile_size
- @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"enable_mathdx_gemm": False})
def kernel(
# Data in:
nefc_in: wp.array[int],
@@ -2107,7 +2323,7 @@ def _update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int,
TILE_SIZE_K = tile_size
- @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"enable_mathdx_gemm": False})
def kernel(
# Model:
M_colind: wp.array[int], # column index of each CSR entry
@@ -2178,349 +2394,26 @@ def _update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int,
return kernel
-# TODO(thowell): combine with JTDAJ ?
-@wp.kernel
-def _update_gradient_JTCJ_sparse(
- # Model:
- opt_impratio_invsqrt: wp.array[float],
- # Data in:
- contact_dist_in: wp.array[float],
- contact_includemargin_in: wp.array[float],
- contact_friction_in: wp.array[types.vec5],
- contact_dim_in: wp.array[int],
- contact_efc_address_in: wp.array2d[int],
- contact_worldid_in: wp.array[int],
- efc_J_rownnz_in: wp.array2d[int],
- efc_J_rowadr_in: wp.array2d[int],
- efc_J_colind_in: wp.array3d[int],
- efc_J_in: wp.array3d[float],
- efc_D_in: wp.array2d[float],
- efc_state_in: wp.array2d[int],
- naconmax_in: int,
- nacon_in: wp.array[int],
+@wp.func
+def _elliptic_hessian_entry_from_projections(
# In:
- ctx_Jaref_in: wp.array2d[float],
- ctx_done_in: wp.array[bool],
- nblocks_perblock: int,
- dim_block: int,
- # Out:
- ctx_h_out: wp.array3d[float],
-):
- conid_start, pairid = wp.tid()
-
- for i in range(nblocks_perblock):
- conid = conid_start + i * dim_block
-
- if conid >= min(nacon_in[0], naconmax_in):
- return
-
- worldid = contact_worldid_in[conid]
- if ctx_done_in[worldid]:
- continue
-
- condim = contact_dim_in[conid]
-
- if condim == 1:
- continue
-
- # check contact status
- if contact_dist_in[conid] - contact_includemargin_in[conid] >= 0.0:
- continue
-
- efcid0 = contact_efc_address_in[conid, 0]
- if efcid0 < 0:
- continue
- if efc_state_in[worldid, efcid0] != types.ConstraintState.CONE:
- continue
-
- # One thread per (contact, support-pair): the support dofs are exactly the colind entries,
- # so decode pairid -> (pos1, pos2) with pos1 <= pos2 directly. No colind scan, and no
- # membership skip (which the all-dof-pairs version wasted on ~99% absent dofs).
- rownnz = efc_J_rownnz_in[worldid, efcid0]
- npairs = rownnz * (rownnz + 1) // 2
- if pairid >= npairs:
- continue
- rowadr0 = efc_J_rowadr_in[worldid, efcid0]
- pos1 = int(0)
- rem = pairid
- while rem >= rownnz - pos1:
- rem -= rownnz - pos1
- pos1 += 1
- pos2 = pos1 + rem
- dofa = efc_J_colind_in[worldid, 0, rowadr0 + pos1]
- dofb = efc_J_colind_in[worldid, 0, rowadr0 + pos2]
- dof1id = wp.min(dofa, dofb)
- dof2id = wp.max(dofa, dofb)
-
- fri = contact_friction_in[conid]
- mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-
- mu2 = mu * mu
- dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2))
-
- if dm == 0.0:
- continue
-
- n = ctx_Jaref_in[worldid, efcid0] * mu
- u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0)
-
- tt = float(0.0)
- for j in range(1, condim):
- efcidj = contact_efc_address_in[conid, j]
- if efcidj >= 0:
- uj = ctx_Jaref_in[worldid, efcidj] * fri[j - 1]
- else:
- uj = 0.0
- tt += uj * uj
- u[j] = uj
-
- if tt <= 0.0:
- t = 0.0
- else:
- t = wp.sqrt(tt)
- t = wp.max(t, types.MJ_MINVAL)
- ttt = wp.max(t * t * t, types.MJ_MINVAL)
-
- # Precompute common subexpressions.
- mu_over_t = math.safe_div(mu, t)
- mu_n_over_ttt = mu * math.safe_div(n, ttt)
- mu2_minus_mu_n_over_t = mu2 - mu * math.safe_div(n, t)
-
- h = float(0.0)
-
- for dim1id in range(condim):
- if dim1id == 0:
- rowadr1 = rowadr0
- dm_fri1 = dm * mu
- else:
- efcid1 = contact_efc_address_in[conid, dim1id]
- if efcid1 < 0:
- continue
- rowadr1 = efc_J_rowadr_in[worldid, efcid1]
- dm_fri1 = dm * fri[dim1id - 1]
-
- # Direct J reads using cached sparse positions.
- efc_J11 = efc_J_in[worldid, 0, rowadr1 + pos1]
- efc_J12 = efc_J_in[worldid, 0, rowadr1 + pos2]
-
- ui = u[dim1id]
-
- for dim2id in range(0, dim1id + 1):
- if dim2id == 0:
- rowadr2 = rowadr0
- dm_fri12 = dm_fri1 * mu
- else:
- efcid2 = contact_efc_address_in[conid, dim2id]
- if efcid2 < 0:
- continue
- rowadr2 = efc_J_rowadr_in[worldid, efcid2]
- dm_fri12 = dm_fri1 * fri[dim2id - 1]
-
- # Direct J reads using cached sparse positions.
- efc_J21 = efc_J_in[worldid, 0, rowadr2 + pos1]
- efc_J22 = efc_J_in[worldid, 0, rowadr2 + pos2]
-
- uj = u[dim2id]
-
- # set first row/column: (1, -mu/t * u)
- if dim1id == 0 and dim2id == 0:
- hcone = 1.0
- elif dim1id == 0:
- hcone = -mu_over_t * uj
- elif dim2id == 0:
- hcone = -mu_over_t * ui
- else:
- hcone = mu_n_over_ttt * ui * uj
-
- # add to diagonal: mu^2 - mu * n / t
- if dim1id == dim2id:
- hcone += mu2_minus_mu_n_over_t
-
- hcone *= dm_fri12
-
- if hcone != 0.0:
- h += hcone * efc_J11 * efc_J22
-
- if dim1id != dim2id:
- h += hcone * efc_J12 * efc_J21
-
- # multiple contacts can contribute to the same (dof1id, dof2id); atomic_add is exact
- wp.atomic_add(ctx_h_out[worldid, dof1id], dof2id, h)
-
-
-@wp.kernel
-def _update_gradient_JTCJ_compact(
- # Model:
- opt_impratio_invsqrt: wp.array[float],
- # Data in:
- contact_dist_in: wp.array[float],
- contact_includemargin_in: wp.array[float],
- contact_friction_in: wp.array[types.vec5],
- contact_dim_in: wp.array[int],
- contact_efc_address_in: wp.array2d[int],
- contact_worldid_in: wp.array[int],
- efc_J_rownnz_in: wp.array2d[int],
- efc_J_rowadr_in: wp.array2d[int],
- efc_J_colind_in: wp.array3d[int],
- efc_J_in: wp.array3d[float],
- efc_D_in: wp.array2d[float],
- efc_state_in: wp.array2d[int],
- dof_cdof_in: wp.array2d[int],
- naconmax_in: int,
- nacon_in: wp.array[int],
- # In:
- ctx_Jaref_in: wp.array2d[float],
- ctx_done_in: wp.array[bool],
- nblocks_perblock: int,
- dim_block: int,
- # Out:
- ctx_h_out: wp.array3d[float],
-):
- conid_start, pairid = wp.tid()
-
- for i in range(nblocks_perblock):
- conid = conid_start + i * dim_block
-
- if conid >= min(nacon_in[0], naconmax_in):
- return
-
- worldid = contact_worldid_in[conid]
- if ctx_done_in[worldid]:
- continue
-
- condim = contact_dim_in[conid]
-
- if condim == 1:
- continue
-
- # check contact status
- if contact_dist_in[conid] - contact_includemargin_in[conid] >= 0.0:
- continue
-
- efcid0 = contact_efc_address_in[conid, 0]
- if efcid0 < 0:
- continue
- if efc_state_in[worldid, efcid0] != types.ConstraintState.CONE:
- continue
-
- rownnz = efc_J_rownnz_in[worldid, efcid0]
- npairs = rownnz * (rownnz + 1) // 2
- if pairid >= npairs:
- continue
-
- rowadr0 = efc_J_rowadr_in[worldid, efcid0]
- pos1 = int(0)
- rem = pairid
- while rem >= rownnz - pos1:
- rem -= rownnz - pos1
- pos1 += 1
- pos2 = pos1 + rem
-
- dofa = efc_J_colind_in[worldid, 0, rowadr0 + pos1]
- dofb = efc_J_colind_in[worldid, 0, rowadr0 + pos2]
-
- # Map to compacted DOFs
- dof1id = dof_cdof_in[worldid, dofa]
- dof2id = dof_cdof_in[worldid, dofb]
-
- if dof1id < 0 or dof2id < 0:
- continue
-
- c_dof1 = wp.min(dof1id, dof2id)
- c_dof2 = wp.max(dof1id, dof2id)
-
- fri = contact_friction_in[conid]
- mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-
- mu2 = mu * mu
- dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2))
-
- if dm == 0.0:
- continue
-
- n = ctx_Jaref_in[worldid, efcid0] * mu
- u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0)
-
- tt = float(0.0)
- for j in range(1, condim):
- efcidj = contact_efc_address_in[conid, j]
- if efcidj >= 0:
- uj = ctx_Jaref_in[worldid, efcidj] * fri[j - 1]
- else:
- uj = 0.0
- tt += uj * uj
- u[j] = uj
-
- if tt <= 0.0:
- t = 0.0
- else:
- t = wp.sqrt(tt)
- t = wp.max(t, types.MJ_MINVAL)
- ttt = wp.max(t * t * t, types.MJ_MINVAL)
-
- # Precompute common subexpressions.
- mu_over_t = math.safe_div(mu, t)
- mu_n_over_ttt = mu * math.safe_div(n, ttt)
- mu2_minus_mu_n_over_t = mu2 - mu * math.safe_div(n, t)
-
- h = float(0.0)
-
- for dim1id in range(condim):
- if dim1id == 0:
- efcid1 = efcid0
- dm_fri1 = dm * mu
- else:
- efcid1 = contact_efc_address_in[conid, dim1id]
- if efcid1 < 0:
- continue
- dm_fri1 = dm * fri[dim1id - 1]
-
- # Read from the compacted dense Jacobian (efc_J_in) using the mapped compacted DOFs
- efc_J11 = efc_J_in[worldid, efcid1, c_dof1]
- efc_J12 = efc_J_in[worldid, efcid1, c_dof2]
-
- ui = u[dim1id]
-
- for dim2id in range(0, dim1id + 1):
- if dim2id == 0:
- efcid2 = efcid0
- dm_fri12 = dm_fri1 * mu
- else:
- efcid2 = contact_efc_address_in[conid, dim2id]
- if efcid2 < 0:
- continue
- dm_fri12 = dm_fri1 * fri[dim2id - 1]
-
- # Read from the compacted dense Jacobian using the mapped compacted DOFs
- efc_J21 = efc_J_in[worldid, efcid2, c_dof1]
- efc_J22 = efc_J_in[worldid, efcid2, c_dof2]
-
- uj = u[dim2id]
-
- # set first row/column: (1, -mu/t * u)
- if dim1id == 0 and dim2id == 0:
- hcone = 1.0
- elif dim1id == 0:
- hcone = -mu_over_t * uj
- elif dim2id == 0:
- hcone = -mu_over_t * ui
- else:
- hcone = mu_n_over_ttt * ui * uj
-
- # add to diagonal: mu^2 - mu * n / t
- if dim1id == dim2id:
- hcone += mu2_minus_mu_n_over_t
-
- hcone *= dm_fri12
-
- if hcone != 0.0:
- h += hcone * efc_J11 * efc_J22
-
- if dim1id != dim2id:
- h += hcone * efc_J12 * efc_J21
-
- # multiple contacts can contribute to the same (c_dof1, c_dof2); atomic_add is exact
- wp.atomic_add(ctx_h_out[worldid, c_dof1], c_dof2, h)
+ dm: float,
+ mu_over_t: float,
+ mu_n_over_ttt: float,
+ tangent_diag: float,
+ z01: float,
+ z02: float,
+ projection1: float,
+ projection2: float,
+ tangent_dot: float,
+) -> float:
+ # Contract the diagonal-plus-rank-one curvature without materializing the cone Hessian.
+ return dm * (
+ z01 * z02
+ - mu_over_t * (z01 * projection2 + z02 * projection1)
+ + mu_n_over_ttt * projection1 * projection2
+ + tangent_diag * tangent_dot
+ )
@wp.kernel
@@ -2589,85 +2482,38 @@ def _update_gradient_JTCJ_dense(
continue
n = ctx_Jaref_in[worldid, efcid0] * mu
- u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0)
-
+ z01 = mu * efc_J_in[worldid, efcid0, dof1id]
+ z02 = mu * efc_J_in[worldid, efcid0, dof2id]
tt = float(0.0)
- for j in range(1, condim):
- efcidj = contact_efc_address_in[conid, j]
- if efcidj >= 0:
- uj = ctx_Jaref_in[worldid, efcidj] * fri[j - 1]
- else:
- uj = 0.0
- tt += uj * uj
- u[j] = uj
+ projection1 = float(0.0)
+ projection2 = float(0.0)
+ tangent_dot = float(0.0)
+ for dim in range(1, condim):
+ efcid = contact_efc_address_in[conid, dim]
+ if efcid >= 0:
+ scale = fri[dim - 1]
+ u = ctx_Jaref_in[worldid, efcid] * scale
+ z1 = scale * efc_J_in[worldid, efcid, dof1id]
+ z2 = scale * efc_J_in[worldid, efcid, dof2id]
+ tt += u * u
+ projection1 += u * z1
+ projection2 += u * z2
+ tangent_dot += z1 * z2
- if tt <= 0.0:
- t = 0.0
- else:
- t = wp.sqrt(tt)
- t = wp.max(t, types.MJ_MINVAL)
+ t = wp.max(wp.sqrt(tt), types.MJ_MINVAL)
ttt = wp.max(t * t * t, types.MJ_MINVAL)
-
- h = float(0.0)
-
- for dim1id in range(condim):
- if dim1id == 0:
- efcid1 = efcid0
- else:
- efcid1 = contact_efc_address_in[conid, dim1id]
- if efcid1 < 0:
- continue
-
- efc_J11 = efc_J_in[worldid, efcid1, dof1id]
- efc_J12 = efc_J_in[worldid, efcid1, dof2id]
-
- ui = u[dim1id]
-
- for dim2id in range(0, dim1id + 1):
- if dim2id == 0:
- efcid2 = efcid0
- else:
- efcid2 = contact_efc_address_in[conid, dim2id]
- if efcid2 < 0:
- continue
-
- efc_J21 = efc_J_in[worldid, efcid2, dof1id]
- efc_J22 = efc_J_in[worldid, efcid2, dof2id]
-
- uj = u[dim2id]
-
- # set first row/column: (1, -mu/t * u)
- if dim1id == 0 and dim2id == 0:
- hcone = 1.0
- elif dim1id == 0:
- hcone = -math.safe_div(mu, t) * uj
- elif dim2id == 0:
- hcone = -math.safe_div(mu, t) * ui
- else:
- hcone = mu * math.safe_div(n, ttt) * ui * uj
-
- # add to diagonal: mu^2 - mu * n / t
- if dim1id == dim2id:
- hcone += mu2 - mu * math.safe_div(n, t)
-
- # pre and post multiply by diag(mu, friction) scale by dm
- if dim1id == 0:
- fri1 = mu
- else:
- fri1 = fri[dim1id - 1]
-
- if dim2id == 0:
- fri2 = mu
- else:
- fri2 = fri[dim2id - 1]
-
- hcone *= dm * fri1 * fri2
-
- if hcone != 0.0:
- h += hcone * efc_J11 * efc_J22
-
- if dim1id != dim2id:
- h += hcone * efc_J12 * efc_J21
+ mu_tinv = math.safe_div(mu, t)
+ h = _elliptic_hessian_entry_from_projections(
+ dm,
+ mu_tinv,
+ mu * math.safe_div(n, ttt),
+ mu2 - n * mu_tinv,
+ z01,
+ z02,
+ projection1,
+ projection2,
+ tangent_dot,
+ )
ctx_h_out[worldid, dof1id, dof2id] += h
@@ -2676,15 +2522,17 @@ def _update_gradient_JTCJ_dense(
def _update_gradient_cholesky(tile_size: int, skip_noflip: bool = False):
SKIP_NOFLIP = skip_noflip
- @wp.kernel(module="unique", enable_backward=False)
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
def kernel(
# In:
ctx_grad_in: wp.array2d[float],
h_in: wp.array3d[float],
- changed_count_in: wp.array[int],
+ state_changed_count_in: wp.array[int],
ctx_done_in: wp.array[bool],
# Out:
- ctx_Mgrad_out: wp.array2d[float],
+ ctx_search_out: wp.array2d[float],
+ ctx_search_dot_out: wp.array[float],
+ ctx_newton_decrement_out: wp.array[float],
):
worldid = wp.tid()
TILE_SIZE = wp.static(tile_size)
@@ -2694,21 +2542,24 @@ def _update_gradient_cholesky(tile_size: int, skip_noflip: bool = False):
# Fast path: skip the solve (see the blocked skip_unchanged variant).
if wp.static(SKIP_NOFLIP):
- if changed_count_in[worldid] == 0:
+ if state_changed_count_in[worldid] == 0:
return
mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE))
wp.tile_cholesky_inplace(mat_tile, fill_mode="upper")
input_tile = wp.tile_load(ctx_grad_in[worldid], shape=TILE_SIZE)
output_tile = wp.tile_cholesky_solve(mat_tile, input_tile, fill_mode="upper")
- wp.tile_store(ctx_Mgrad_out[worldid], output_tile)
+ sums = wp.tile_reduce(wp.add, wp.tile_map(solve_search_sums, input_tile, output_tile))[0]
+ ctx_search_dot_out[worldid] = sums[0]
+ ctx_newton_decrement_out[worldid] = sums[1]
+ wp.tile_store(ctx_search_out[worldid], wp.tile_map(wp.mul, output_tile, -1.0))
return kernel
@cache_kernel
-def _update_gradient_cholesky_blocked(tile_size: int, matrix_size: int, check_skip: bool = True):
- @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
+def _update_gradient_cholesky_blocked(tile_size: int, matrix_size: int, vector_size: int):
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"enable_mathdx_gemm": False})
def kernel(
# In:
ctx_done_in: wp.array[bool],
@@ -2716,42 +2567,9 @@ def _update_gradient_cholesky_blocked(tile_size: int, matrix_size: int, check_sk
ctx_h_in: wp.array3d[float],
ctx_hfactor: wp.array3d[float],
# Out:
- ctx_Mgrad_out: wp.array3d[float],
- ):
- worldid = wp.tid()
- TILE_SIZE = wp.static(tile_size)
-
- if wp.static(check_skip):
- if ctx_done_in[worldid]:
- return
-
- # We need matrix size both as a runtime input as well as a static input:
- # static input is needed to specify the tile sizes for the compiler
- # runtime input is needed for the loop bounds, otherwise warp will unroll
- # unconditionally leading to shared memory capacity issues.
-
- wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
- ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
- )
-
- return kernel
-
-
-@cache_kernel
-def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int, skip_noflip: bool = False):
- """Blocked Cholesky that skips factorization when no constraints changed."""
- SKIP_NOFLIP = skip_noflip
-
- @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
- def kernel(
- # In:
- ctx_done_in: wp.array[bool],
- ctx_grad_in: wp.array3d[float],
- ctx_h_in: wp.array3d[float],
- changed_count_in: wp.array[int],
- ctx_hfactor: wp.array3d[float],
- # Out:
- ctx_Mgrad_out: wp.array3d[float],
+ ctx_search_out: wp.array3d[float],
+ ctx_search_dot_out: wp.array[float],
+ ctx_newton_decrement_out: wp.array[float],
):
worldid = wp.tid()
TILE_SIZE = wp.static(tile_size)
@@ -2759,29 +2577,105 @@ def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size
if ctx_done_in[worldid]:
return
- # Fast path: skip the solve; Mgrad stays stale on the unchanged ray, and
- # the linesearch is invariant to the scale of its direction.
- if wp.static(SKIP_NOFLIP):
- if changed_count_in[worldid] == 0:
- return
+ # We need matrix size both as a runtime input as well as a static input:
+ # static input is needed to specify the tile sizes for the compiler
+ # runtime input is needed for the loop bounds, otherwise warp will unroll
+ # unconditionally leading to shared memory capacity issues.
- wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
- ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
- )
- else:
- if changed_count_in[worldid] > 0:
- wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
- ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
- )
- else:
- wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
- ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
- )
+ sums = wp.static(create_blocked_cholesky_augmented_factorize_solve_newton_func(TILE_SIZE, matrix_size, vector_size))(
+ ctx_h_in[worldid],
+ ctx_grad_in[worldid],
+ matrix_size,
+ ctx_hfactor[worldid],
+ ctx_search_out[worldid],
+ )
+ ctx_search_dot_out[worldid] = sums[0]
+ ctx_newton_decrement_out[worldid] = sums[1]
return kernel
-@wp.kernel
+@cache_kernel
+def _cholesky_factorize_solve_blocked(tile_size: int, matrix_size: int):
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"enable_mathdx_gemm": False})
+ def kernel(
+ # In:
+ A_in: wp.array3d[float],
+ b_in: wp.array3d[float],
+ # Out:
+ U_out: wp.array3d[float],
+ x_out: wp.array3d[float],
+ ):
+ worldid = wp.tid()
+ TILE_SIZE = wp.static(tile_size)
+
+ wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
+ A_in[worldid],
+ b_in[worldid],
+ matrix_size,
+ U_out[worldid],
+ x_out[worldid],
+ )
+
+ return kernel
+
+
+@cache_kernel
+def _update_gradient_cholesky_blocked_skip_unchanged(
+ tile_size: int, matrix_size: int, vector_size: int, skip_noflip: bool = False
+):
+ """Blocked Cholesky that skips factorization when no constraints changed."""
+ SKIP_NOFLIP = skip_noflip
+
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"enable_mathdx_gemm": False})
+ def kernel(
+ # In:
+ ctx_done_in: wp.array[bool],
+ ctx_grad_in: wp.array3d[float],
+ ctx_h_in: wp.array3d[float],
+ quad_changed_count_in: wp.array[int],
+ state_changed_count_in: wp.array[int],
+ ctx_hfactor: wp.array3d[float],
+ # Out:
+ ctx_search_out: wp.array3d[float],
+ ctx_search_dot_out: wp.array[float],
+ ctx_newton_decrement_out: wp.array[float],
+ ):
+ worldid = wp.tid()
+ TILE_SIZE = wp.static(tile_size)
+
+ if ctx_done_in[worldid]:
+ return
+
+ # The linesearch is invariant to direction scale, so an unchanged ray can
+ # keep its previous search direction.
+ if wp.static(SKIP_NOFLIP):
+ if state_changed_count_in[worldid] == 0:
+ return
+
+ if quad_changed_count_in[worldid] > 0:
+ sums = wp.static(create_blocked_cholesky_augmented_factorize_solve_newton_func(TILE_SIZE, matrix_size, vector_size))(
+ ctx_h_in[worldid],
+ ctx_grad_in[worldid],
+ matrix_size,
+ ctx_hfactor[worldid],
+ ctx_search_out[worldid],
+ )
+ else:
+ sums = wp.static(create_blocked_cholesky_solve_newton_func(TILE_SIZE, matrix_size, vector_size))(
+ ctx_hfactor[worldid],
+ ctx_grad_in[worldid],
+ matrix_size,
+ ctx_search_out[worldid],
+ )
+
+ ctx_search_dot_out[worldid] = sums[0]
+ ctx_newton_decrement_out[worldid] = sums[1]
+
+ return kernel
+
+
+@wp.kernel(grid_stride=True)
def _padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float]):
worldid, elementid = wp.tid()
@@ -2795,7 +2689,7 @@ def _padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float
def _cholesky_factorize_solve(
m: types.Model, d: types.Data, ctx: SolverContext, skip_unchanged: bool = False, skip_noflip: bool = False
):
- """Cholesky factorize ctx.h and solve for Mgrad.
+ """Cholesky factorize ctx.h and form the Newton search direction.
If skip_unchanged is True (blocked path only), worlds where no constraints
changed reuse the cached factorization in hfactor instead of refactorizing.
@@ -2804,8 +2698,8 @@ def _cholesky_factorize_solve(
wp.launch_tiled(
_update_gradient_cholesky(m.nv, skip_noflip),
dim=d.nworld,
- inputs=[ctx.grad, ctx.h, ctx.changed_efc_count if skip_noflip else d.nefc, ctx.done],
- outputs=[ctx.Mgrad],
+ inputs=[ctx.grad, ctx.h, ctx.state_changed_count if skip_noflip else d.nefc, ctx.done],
+ outputs=[ctx.search, ctx.search_dot, ctx.newton_decrement],
block_dim=m.block_dim.update_gradient_cholesky,
)
else:
@@ -2818,90 +2712,389 @@ def _cholesky_factorize_solve(
if skip_unchanged:
wp.launch_tiled(
- _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad, skip_noflip),
+ _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad, m.nv, skip_noflip),
dim=d.nworld,
- inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.changed_efc_count, ctx.hfactor],
- outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))],
+ inputs=[
+ ctx.done,
+ ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)),
+ ctx.h,
+ ctx.quad_changed_count,
+ ctx.state_changed_count if skip_noflip else ctx.quad_changed_count,
+ ctx.hfactor,
+ ],
+ outputs=[
+ ctx.search.reshape(shape=(d.nworld, m.nv, 1)),
+ ctx.search_dot,
+ ctx.newton_decrement,
+ ],
block_dim=m.block_dim.update_gradient_cholesky_blocked,
)
else:
wp.launch_tiled(
- _update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
+ _update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad, m.nv),
dim=d.nworld,
inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.hfactor],
- outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))],
+ outputs=[
+ ctx.search.reshape(shape=(d.nworld, m.nv, 1)),
+ ctx.search_dot,
+ ctx.newton_decrement,
+ ],
block_dim=m.block_dim.update_gradient_cholesky_blocked,
)
# ---------------------------------------------------------------------------
-# H += J^T D J. D diagonal, so each efc row adds one rank-1 outer product. make_constraint
-# groups a constraint's contiguous efc rows (shared colind = dof support S) into one |S|x|S|
-# block, stored densely per world in efc.jtdaj_{adr,nrow,nblock}. The launch fills the
-# GPU once (groups_per_world slots/world) then grid-strides the rest, so no thread lands on a
-# non-head efc row. A block's upper-triangular entries split across THREADS_PER_GROUP threads
-# (one warp -> coalesced J reads); entry -> (block_row, block_col) is the triangular-number
-# inverse, exact in float32 since column boundaries are perfect squares (8*entry+1 = (2c+1)^2).
+# Constraint groups contain consecutive rows with identical sparse support. Each thread group
+# accumulates their upper-triangular Hessian block, including elliptic cone curvature.
# ---------------------------------------------------------------------------
-_JTDAJ_THREADS_PER_GROUP = 32 # one warp per group, so its J reads coalesce
-_JTDAJ_OVERSUBSCRIBE_WAVES = 6 # grid-stride depth; short per-warp chains load-balance groups
+_JTDAJ_THREADS_PER_GROUP = 32
+_JTDAJ_OVERSUBSCRIBE_WAVES = 6
+
+
+@cache_kernel
+def _JTDACJ_sparse(compact: bool, cone_type: types.ConeType, max_condim: int):
+ COMPACT = compact
+ ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
+ MAX_CONDIM = max_condim
+
+ def make_curvature_terms(condim: int):
+ @wp.func
+ def func(
+ # Model:
+ opt_impratio_invsqrt: wp.array[float],
+ # Data in:
+ efc_D_in: wp.array2d[float],
+ # In:
+ fri: types.vec5,
+ ctx_Jaref_in: wp.array2d[float],
+ worldid: int,
+ efcid0: int,
+ block_rows: int,
+ ) -> types.vec16:
+ mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+ mu2 = mu * mu
+ dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2))
+ if dm == 0.0:
+ return types.vec16()
+
+ n = ctx_Jaref_in[worldid, efcid0] * mu
+ terms = types.vec16()
+ terms[6] = mu
+ tt = float(0.0)
+ for dim in range(1, wp.static(condim)):
+ if dim < block_rows:
+ efcid = efcid0 + dim
+ scale = fri[dim - 1]
+ u = ctx_Jaref_in[worldid, efcid] * scale
+ terms[dim] = u
+ terms[6 + dim] = scale
+ tt += u * u
+
+ t = wp.max(wp.sqrt(tt), types.MJ_MINVAL)
+ ttt = wp.max(t * t * t, types.MJ_MINVAL)
+ mu_over_t = math.safe_div(mu, t)
+ mu_n_over_ttt = mu * math.safe_div(n, ttt)
+ tangent_diag = mu2 - n * mu_over_t
+
+ # Layout: tangent u[1:6], scales[6:12], dm, mu/t, mu*n/t^3, tangent diagonal.
+ terms[12] = dm
+ terms[13] = mu_over_t
+ terms[14] = mu_n_over_ttt
+ terms[15] = tangent_diag
+ return terms
+
+ return func
+
+ def make_hessian_entry(condim: int):
+ @wp.func
+ def func(
+ # Data in:
+ efc_J_in: wp.array3d[float],
+ # In:
+ terms: Any,
+ rowadr: types.vec6i,
+ worldid: int,
+ pos1: int,
+ pos2: int,
+ ) -> float:
+ z01 = terms[6] * efc_J_in[worldid, 0, rowadr[0] + pos1]
+ z02 = terms[6] * efc_J_in[worldid, 0, rowadr[0] + pos2]
+ projection1 = float(0.0)
+ projection2 = float(0.0)
+ tangent_dot = float(0.0)
+ for dim in range(1, wp.static(condim)):
+ z1 = terms[6 + dim] * efc_J_in[worldid, 0, rowadr[dim] + pos1]
+ z2 = terms[6 + dim] * efc_J_in[worldid, 0, rowadr[dim] + pos2]
+ projection1 += terms[dim] * z1
+ projection2 += terms[dim] * z2
+ tangent_dot += z1 * z2
+
+ return _elliptic_hessian_entry_from_projections(
+ terms[12],
+ terms[13],
+ terms[14],
+ terms[15],
+ z01,
+ z02,
+ projection1,
+ projection2,
+ tangent_dot,
+ )
+
+ return func
+
+ curvature_terms3 = make_curvature_terms(3)
+ curvature_terms4 = make_curvature_terms(4)
+ curvature_terms6 = make_curvature_terms(6)
+ hessian_entry3 = make_hessian_entry(3)
+ hessian_entry4 = make_hessian_entry(4)
+ hessian_entry6 = make_hessian_entry(6)
+
+ @wp.func
+ def curvature_terms(
+ # Model:
+ opt_impratio_invsqrt: wp.array[float],
+ # Data in:
+ efc_D_in: wp.array2d[float],
+ # In:
+ fri: types.vec5,
+ ctx_Jaref_in: wp.array2d[float],
+ worldid: int,
+ efcid0: int,
+ condim: int,
+ block_rows: int,
+ ) -> types.vec16:
+ if wp.static(MAX_CONDIM == 3):
+ return curvature_terms3(opt_impratio_invsqrt, efc_D_in, fri, ctx_Jaref_in, worldid, efcid0, block_rows)
+
+ if condim == 3:
+ return curvature_terms3(opt_impratio_invsqrt, efc_D_in, fri, ctx_Jaref_in, worldid, efcid0, block_rows)
+ if wp.static(MAX_CONDIM == 4):
+ return curvature_terms4(opt_impratio_invsqrt, efc_D_in, fri, ctx_Jaref_in, worldid, efcid0, block_rows)
+ if condim == 4:
+ return curvature_terms4(opt_impratio_invsqrt, efc_D_in, fri, ctx_Jaref_in, worldid, efcid0, block_rows)
+ return curvature_terms6(opt_impratio_invsqrt, efc_D_in, fri, ctx_Jaref_in, worldid, efcid0, block_rows)
+
+ @wp.func
+ def hessian_entry(
+ # Data in:
+ efc_J_in: wp.array3d[float],
+ # In:
+ terms: Any,
+ rowadr: types.vec6i,
+ worldid: int,
+ pos1: int,
+ pos2: int,
+ condim: int,
+ ) -> float:
+ if wp.static(MAX_CONDIM == 3):
+ return hessian_entry3(efc_J_in, terms, rowadr, worldid, pos1, pos2)
+
+ if condim == 3:
+ return hessian_entry3(efc_J_in, terms, rowadr, worldid, pos1, pos2)
+ if wp.static(MAX_CONDIM == 4):
+ return hessian_entry4(efc_J_in, terms, rowadr, worldid, pos1, pos2)
+ if condim == 4:
+ return hessian_entry4(efc_J_in, terms, rowadr, worldid, pos1, pos2)
+ return hessian_entry6(efc_J_in, terms, rowadr, worldid, pos1, pos2)
+
+ @wp.kernel(module="unique", enable_backward=False, grid_stride=True)
+ def kernel(
+ # Model:
+ opt_impratio_invsqrt: wp.array[float],
+ # Data in:
+ contact_friction_in: wp.array[types.vec5],
+ contact_dim_in: wp.array[int],
+ efc_id_in: wp.array2d[int],
+ efc_jtdaj_adr_in: wp.array2d[int],
+ efc_jtdaj_nrow_in: wp.array2d[int],
+ efc_jtdaj_nblock_in: wp.array[int],
+ efc_J_rownnz_in: wp.array2d[int],
+ efc_J_rowadr_in: wp.array2d[int],
+ efc_J_colind_in: wp.array3d[int],
+ efc_J_in: wp.array3d[float],
+ efc_D_in: wp.array2d[float],
+ efc_state_in: wp.array2d[int],
+ dof_cdof_in: wp.array2d[int],
+ # In:
+ ctx_Jaref_in: wp.array2d[float],
+ ctx_done_in: wp.array[bool],
+ groups_per_world: int,
+ # Out:
+ h_out: wp.array3d[float],
+ ):
+ worldid, slot, lane = wp.tid()
+ if wp.static(ELLIPTIC):
+ lanes = wp.block_dim()
+ else:
+ lanes = wp.static(_JTDAJ_THREADS_PER_GROUP)
+ if ctx_done_in[worldid]:
+ return
+ count = efc_jtdaj_nblock_in[worldid]
+ for groupid in range(slot, count, groups_per_world):
+ head_row = efc_jtdaj_adr_in[worldid, groupid]
+ block_rows = efc_jtdaj_nrow_in[worldid, groupid]
+ head_adr = efc_J_rowadr_in[worldid, head_row]
+ support = efc_J_rownnz_in[worldid, head_row]
+ n_entries = support * (support + 1) // 2
+
+ is_cone = False
+ if wp.static(ELLIPTIC):
+ is_cone = efc_state_in[worldid, head_row] == types.ConstraintState.CONE
+ condim = int(0)
+ # Clipped cone rows retain the safe head address and a zero scale.
+ cone_rowadr = types.vec6i(head_adr, head_adr, head_adr, head_adr, head_adr, head_adr)
+ if is_cone:
+ conid = efc_id_in[worldid, head_row]
+ if wp.static(MAX_CONDIM == 3):
+ condim = int(3)
+ else:
+ condim = contact_dim_in[conid]
+ for dim in range(1, wp.static(MAX_CONDIM)):
+ if dim < block_rows:
+ cone_rowadr[dim] = head_adr + dim * support
+ local_terms = types.vec16()
+ if lane == 0:
+ local_terms = curvature_terms(
+ opt_impratio_invsqrt,
+ efc_D_in,
+ contact_friction_in[conid],
+ ctx_Jaref_in,
+ worldid,
+ head_row,
+ condim,
+ block_rows,
+ )
+ cone_terms = wp.tile_zeros(shape=(16,), dtype=float, storage="shared")
+ for term_id in range(1, 16):
+ wp.tile_scatter_masked(cone_terms, term_id, local_terms[term_id], lane == 0)
+
+ for entry in range(lane, n_entries, lanes):
+ block_col = int((wp.sqrt(float(8 * entry + 1)) - 1.0) * 0.5)
+ block_row = entry - block_col * (block_col + 1) // 2
+ dof_row = efc_J_colind_in[worldid, 0, head_adr + block_row]
+ dof_col = efc_J_colind_in[worldid, 0, head_adr + block_col]
+ hval = float(0.0)
+ if wp.static(ELLIPTIC):
+ if is_cone:
+ hval = hessian_entry(
+ efc_J_in,
+ cone_terms,
+ cone_rowadr,
+ worldid,
+ block_row,
+ block_col,
+ condim,
+ )
+ if not is_cone:
+ for member in range(block_rows):
+ member_row = head_row + member
+ if efc_state_in[worldid, member_row] == types.ConstraintState.QUADRATIC.value:
+ member_adr = efc_J_rowadr_in[worldid, member_row]
+ j_row = efc_J_in[worldid, 0, member_adr + block_row]
+ j_col = efc_J_in[worldid, 0, member_adr + block_col]
+ hval += j_row * efc_D_in[worldid, member_row] * j_col
+ if hval != 0.0:
+ if wp.static(COMPACT):
+ dof_row = dof_cdof_in[worldid, dof_row]
+ dof_col = dof_cdof_in[worldid, dof_col]
+ if dof_row < 0 or dof_col < 0:
+ continue
+ wp.atomic_add(h_out[worldid, wp.min(dof_row, dof_col)], wp.max(dof_row, dof_col), hval)
+
+ return kernel
+
+
+def _jtdaj_groups_per_world(nworld: int, njmax: int) -> int:
+ # njmax is capacity and often mostly empty, so cap slots at a few resident waves.
+ block_size, min_grid_size = wp.get_suggested_block_size(_JTDACJ_sparse(False, types.ConeType.PYRAMIDAL, 3))
+ device_warps = max(1, block_size * min_grid_size // _JTDAJ_THREADS_PER_GROUP)
+ return max(1, min(njmax, _JTDAJ_OVERSUBSCRIBE_WAVES * device_warps // nworld))
@wp.kernel
-def _JTDAJ_sparse(
+def _diag_precond_build(
+ # Model:
+ body_simple: wp.array[int],
+ dof_bodyid: wp.array[int],
+ M_rownnz: wp.array[int],
+ M_rowadr: wp.array[int],
# Data in:
- efc_jtdaj_adr_in: wp.array2d[int],
- efc_jtdaj_nrow_in: wp.array2d[int],
- efc_jtdaj_nblock_in: wp.array[int],
+ M_in: wp.array2d[float],
+ # In:
+ ctx_done_in: wp.array[bool],
+ # Out:
+ diag_out: wp.array2d[float],
+):
+ """Initialize diagonal with M_ii + regularization for flex DOFs only."""
+ worldid, dofid = wp.tid()
+ if ctx_done_in[worldid]:
+ return
+ if body_simple[dof_bodyid[dofid]] != 2:
+ return
+ madr_ii = M_rowadr[dofid] + M_rownnz[dofid] - 1
+ diag_out[worldid, dofid] = M_in[worldid, madr_ii] + float(1.0e-12)
+
+
+@wp.kernel
+def _diag_precond_add_JTDJ(
+ # Model:
+ body_simple: wp.array[int],
+ dof_bodyid: wp.array[int],
+ # Data in:
+ nefc_in: wp.array[int],
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
efc_J_in: wp.array3d[float],
efc_D_in: wp.array2d[float],
- efc_state_in: wp.array2d[int],
# In:
ctx_done_in: wp.array[bool],
- groups_per_world: int,
# Out:
- h_out: wp.array3d[float],
+ diag_out: wp.array2d[float],
):
- worldid, slot, lane = wp.tid()
+ """Add diagonal of J^T D J for flex DOFs only."""
+ worldid, efcid = wp.tid()
if ctx_done_in[worldid]:
return
- count = efc_jtdaj_nblock_in[worldid]
- for groupid in range(slot, count, groups_per_world): # grid-stride this world's group list
- head_row = efc_jtdaj_adr_in[worldid, groupid]
- block_rows = efc_jtdaj_nrow_in[worldid, groupid]
- head_adr = efc_J_rowadr_in[worldid, head_row]
- support = efc_J_rownnz_in[worldid, head_row] # dofs the constraint touches = block dimension
- n_entries = support * (support + 1) // 2 # upper-triangular entries of the |S|x|S| block
- for entry in range(lane, n_entries, wp.static(_JTDAJ_THREADS_PER_GROUP)):
- block_col = int((wp.sqrt(float(8 * entry + 1)) - 1.0) * 0.5)
- block_row = entry - block_col * (block_col + 1) // 2
- dof_row = efc_J_colind_in[worldid, 0, head_adr + block_row]
- dof_col = efc_J_colind_in[worldid, 0, head_adr + block_col]
- hval = float(0.0)
- for member in range(block_rows):
- member_row = head_row + member
- if efc_state_in[worldid, member_row] == types.ConstraintState.QUADRATIC.value:
- member_adr = efc_J_rowadr_in[worldid, member_row]
- j_row = efc_J_in[worldid, 0, member_adr + block_row]
- j_col = efc_J_in[worldid, 0, member_adr + block_col]
- hval += j_row * efc_D_in[worldid, member_row] * j_col
- if hval != 0.0: # skip the atomic when no member row is active
- wp.atomic_add(h_out[worldid, wp.min(dof_row, dof_col)], wp.max(dof_row, dof_col), hval)
+ if efcid >= nefc_in[worldid]:
+ return
+ D = efc_D_in[worldid, efcid]
+ if D == 0.0:
+ return
+ rownnz = efc_J_rownnz_in[worldid, efcid]
+ rowadr = efc_J_rowadr_in[worldid, efcid]
+ for i in range(rownnz):
+ col = efc_J_colind_in[worldid, 0, rowadr + i]
+ if body_simple[dof_bodyid[col]] != 2:
+ continue
+ Jval = efc_J_in[worldid, 0, rowadr + i]
+ if Jval != 0.0:
+ wp.atomic_add(diag_out, worldid, col, D * Jval * Jval)
-def _jtdaj_groups_per_world(nworld: int, njmax: int) -> int:
- # Per-world width of the grid stride. Target one warp per group-slot (njmax), but cap the grid at
- # _JTDAJ_OVERSUBSCRIBE_WAVES device waves -- else high-njmax worlds dispatch many idle tail warps
- # (njmax >> actual groups). A few waves of oversubscription keep each warp's serial chain short,
- # load-balancing the variable group sizes (measured plateau: ~4-8 waves).
- block_size, min_grid_size = wp.get_suggested_block_size(_JTDAJ_sparse)
- # block_size * min_grid_size = full-device thread count (block_size cancels): the kernel's max
- # resident threads (one wave), a device property independent of nworld and our launch block_dim.
- device_warps = max(1, block_size * min_grid_size // _JTDAJ_THREADS_PER_GROUP)
- return max(1, min(njmax, _JTDAJ_OVERSUBSCRIBE_WAVES * device_warps // nworld))
+@wp.kernel
+def _diag_precond_apply(
+ # Model:
+ body_simple: wp.array[int],
+ dof_bodyid: wp.array[int],
+ # Data in:
+ qLDiagInv_in: wp.array2d[float],
+ # In:
+ diag_in: wp.array2d[float],
+ grad_in: wp.array2d[float],
+ ctx_done_in: wp.array[bool],
+ # Out:
+ Mgrad_out: wp.array2d[float],
+):
+ """Apply preconditioner: flex DOFs use diag(M+JTDJ), others use M diagonal."""
+ worldid, dofid = wp.tid()
+ if ctx_done_in[worldid]:
+ return
+ if body_simple[dof_bodyid[dofid]] == 2:
+ Mgrad_out[worldid, dofid] = grad_in[worldid, dofid] / diag_in[worldid, dofid]
+ else:
+ Mgrad_out[worldid, dofid] = qLDiagInv_in[worldid, dofid] * grad_in[worldid, dofid]
def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact: bool = False):
@@ -2919,7 +3112,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
_update_gradient_zero_grad_dot(False),
dim=d.nworld,
inputs=[d.nefc, ctx.alpha, ctx.done],
- outputs=[ctx.grad_dot, ctx.grad_scale],
+ outputs=[ctx.grad_dot, ctx.newton_decrement, ctx.grad_scale, ctx.search_unchanged],
)
wp.launch(
_update_gradient_grad(False),
@@ -2929,36 +3122,61 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
)
if m.opt.solver == types.SolverType.CG:
- smooth.solve_m(m, d, ctx.Mgrad, ctx.grad)
+ if m.is_sparse and m.nflex > 0:
+ wp.launch(
+ _diag_precond_apply,
+ dim=(d.nworld, m.nv),
+ inputs=[m.body_simple, m.dof_bodyid, d.qLDiagInv, ctx.diag_precond, ctx.grad, ctx.done],
+ outputs=[ctx.Mgrad],
+ )
+ else:
+ smooth.solve_m(m, d, ctx.Mgrad, ctx.grad)
elif m.opt.solver == types.SolverType.NEWTON:
# h = M + (efc_J.T * efc_D * active) @ efc_J
- if m.is_sparse:
+ sc = _sparse_compact(ctx)
+ if m.is_sparse or sc:
+ mj = ctx.compact_m_full if sc else m
+ dj = ctx.compact_d_full if sc else d
wp.launch(
- _update_gradient_init_h_sparse,
+ _update_gradient_init_h_sparse(sc),
dim=(d.nworld, m.nv_pad, m.nv_pad),
- inputs=[m.nv, m.M_elemid, d.M, ctx.done],
+ inputs=[mj.nv, mj.M_elemid, dj.M, dj.cdof_dof, ctx.done],
outputs=[ctx.h],
)
groups_per_world = _jtdaj_groups_per_world(d.nworld, d.njmax)
+ max_condim = 3
+ if m.opt.cone == types.ConeType.ELLIPTIC and m.nmaxcondim > 3:
+ max_condim = int(m.nmaxcondim)
+ jtdaj_kernel = _JTDACJ_sparse(sc, m.opt.cone, max_condim)
+ jtdaj_inputs = [
+ m.opt.impratio_invsqrt,
+ d.contact.friction,
+ d.contact.dim,
+ d.efc.id,
+ dj.efc.jtdaj_adr,
+ dj.efc.jtdaj_nrow,
+ dj.efc.jtdaj_nblock,
+ dj.efc.J_rownnz,
+ dj.efc.J_rowadr,
+ dj.efc.J_colind,
+ dj.efc.J,
+ d.efc.D,
+ d.efc.state,
+ dj.dof_cdof,
+ ctx.Jaref,
+ ctx.done,
+ groups_per_world,
+ ]
+ elliptic = m.opt.cone == types.ConeType.ELLIPTIC
+ threads_per_group = 1 if elliptic and wp.get_device().is_cpu else _JTDAJ_THREADS_PER_GROUP
+ block_dim = threads_per_group if elliptic else mj.block_dim.update_gradient_JTDAJ_sparse
wp.launch(
- _JTDAJ_sparse,
- dim=(d.nworld, groups_per_world, _JTDAJ_THREADS_PER_GROUP),
- inputs=[
- d.efc.jtdaj_adr,
- d.efc.jtdaj_nrow,
- d.efc.jtdaj_nblock,
- d.efc.J_rownnz,
- d.efc.J_rowadr,
- d.efc.J_colind,
- d.efc.J,
- d.efc.D,
- d.efc.state,
- ctx.done,
- groups_per_world,
- ],
+ jtdaj_kernel,
+ dim=(d.nworld, groups_per_world, threads_per_group),
+ inputs=jtdaj_inputs,
outputs=[ctx.h],
- block_dim=m.block_dim.update_gradient_JTDAJ_sparse,
+ block_dim=block_dim,
)
else:
if compact:
@@ -2995,7 +3213,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
block_dim=m.block_dim.update_gradient_JTDAJ_dense,
)
- if m.opt.cone == types.ConeType.ELLIPTIC:
+ if m.opt.cone == types.ConeType.ELLIPTIC and not (m.is_sparse or sc):
# Optimization: launching update_gradient_JTCJ with limited number of blocks on a GPU.
# Profiling suggests that only a fraction of blocks out of the original
# d.njmax blocks do the actual work. It aims to minimize #CTAs with no
@@ -3003,17 +3221,6 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
# of SMs on the GPU. We can now query the SM count:
# https://github.com/NVIDIA/warp/commit/f3814e7e5459e5fd13032cf0fddb3daddd510f30
- # Block-limit the launch: cap the grid near SM-filling width and stride over contacts, so
- # we don't over-launch naconmax (capacity) threads when active contacts are far fewer. The
- # sparse kernel uses one thread per (contact, support-pair) (jtcj_max_pairs), the dense one
- # per (contact, dof-pair) (dof_tri_row.size).
- # `compact` is set by solve_compact's inner solve, which runs the dense factor/solve on
- # the nvmax_pad block but maps sparse contact support-pairs to compacted DOFs via dof_cdof.
- # (Don't infer it from `d.nvmax < m.nv`: after solve_compact's shallow m2/d2 replace that
- # reduces to `nvmax < nvmax_pad`, which is false whenever nvmax is a tile multiple and
- # silently falls back to the O(nvmax_pad^2) dense cone scan.)
- is_sparse_compact = compact and (d.efc.J_colind.shape[1] > 0)
- jtcj_second_dim = m.jtcj_max_pairs if (m.is_sparse or is_sparse_compact) else m.dof_tri_row.size
if wp.get_device().is_cuda:
sm_count = wp.get_device().sm_count
@@ -3021,95 +3228,38 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
# can be changed in the future to fine-tune the perf. The optimal factor will
# depend on the kernel's occupancy, which determines how many blocks can
# simultaneously run on the SM. TODO: This factor can be tuned further.
- dim_block = ceil((sm_count * 6 * 256) / jtcj_second_dim)
+ dim_block = ceil((sm_count * 6 * 256) / m.dof_tri_row.size)
else:
# fall back for CPU
dim_block = d.naconmax
nblocks_perblock = int((d.naconmax + dim_block - 1) / dim_block)
- if m.is_sparse:
- wp.launch(
- _update_gradient_JTCJ_sparse,
- dim=(dim_block, m.jtcj_max_pairs),
- inputs=[
- m.opt.impratio_invsqrt,
- d.contact.dist,
- d.contact.includemargin,
- d.contact.friction,
- d.contact.dim,
- d.contact.efc_address,
- d.contact.worldid,
- d.efc.J_rownnz,
- d.efc.J_rowadr,
- d.efc.J_colind,
- d.efc.J,
- d.efc.D,
- d.efc.state,
- d.naconmax,
- d.nacon,
- ctx.Jaref,
- ctx.done,
- nblocks_perblock,
- dim_block,
- ],
- outputs=[ctx.h],
- )
- else:
- if is_sparse_compact:
- wp.launch(
- _update_gradient_JTCJ_compact,
- dim=(dim_block, m.jtcj_max_pairs),
- inputs=[
- m.opt.impratio_invsqrt,
- d.contact.dist,
- d.contact.includemargin,
- d.contact.friction,
- d.contact.dim,
- d.contact.efc_address,
- d.contact.worldid,
- d.efc.J_rownnz,
- d.efc.J_rowadr,
- d.efc.J_colind,
- d.efc.J,
- d.efc.D,
- d.efc.state,
- d.dof_cdof,
- d.naconmax,
- d.nacon,
- ctx.Jaref,
- ctx.done,
- nblocks_perblock,
- dim_block,
- ],
- outputs=[ctx.h],
- )
- else:
- wp.launch(
- _update_gradient_JTCJ_dense,
- dim=(dim_block, m.dof_tri_row.size),
- inputs=[
- m.opt.impratio_invsqrt,
- m.dof_tri_row,
- m.dof_tri_col,
- d.contact.dist,
- d.contact.includemargin,
- d.contact.friction,
- d.contact.dim,
- d.contact.efc_address,
- d.contact.worldid,
- d.efc.J,
- d.efc.D,
- d.efc.state,
- d.naconmax,
- d.nacon,
- ctx.Jaref,
- ctx.done,
- nblocks_perblock,
- dim_block,
- ],
- outputs=[ctx.h],
- )
+ wp.launch(
+ _update_gradient_JTCJ_dense,
+ dim=(dim_block, m.dof_tri_row.size),
+ inputs=[
+ m.opt.impratio_invsqrt,
+ m.dof_tri_row,
+ m.dof_tri_col,
+ d.contact.dist,
+ d.contact.includemargin,
+ d.contact.friction,
+ d.contact.dim,
+ d.contact.efc_address,
+ d.contact.worldid,
+ d.efc.J,
+ d.efc.D,
+ d.efc.state,
+ d.naconmax,
+ d.nacon,
+ ctx.Jaref,
+ ctx.done,
+ nblocks_perblock,
+ dim_block,
+ ],
+ outputs=[ctx.h],
+ )
_cholesky_factorize_solve(m, d, ctx)
else:
@@ -3122,12 +3272,12 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
Skips the full J^T*D*J rebuild by applying only the delta from constraints
that changed QUADRATIC state, then re-factorizes and solves.
"""
- changed = ctx.changed_efc_count if stable_fast else d.nefc
+ changed = ctx.state_changed_count if stable_fast else d.nefc
wp.launch(
_update_gradient_zero_grad_dot(stable_fast),
dim=d.nworld,
inputs=[changed, ctx.alpha, ctx.done],
- outputs=[ctx.grad_dot, ctx.grad_scale],
+ outputs=[ctx.grad_dot, ctx.newton_decrement, ctx.grad_scale, ctx.search_unchanged],
)
wp.launch(
@@ -3138,20 +3288,23 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
)
# Update upper triangle of H with delta from changed constraints.
- if m.is_sparse:
- slots = _jtdaj_groups_per_world(d.nworld, ctx.changed_efc_ids.shape[1])
+ sc = _sparse_compact(ctx)
+ if m.is_sparse or sc:
+ dj = ctx.compact_d_full if sc else d
+ slots = _jtdaj_groups_per_world(d.nworld, ctx.quad_changed_ids.shape[1])
wp.launch(
- _update_gradient_h_incremental_sparse,
+ _update_gradient_h_incremental_sparse(sc),
dim=(d.nworld, slots, _JTDAJ_THREADS_PER_GROUP),
inputs=[
- d.efc.J_rownnz,
- d.efc.J_rowadr,
- d.efc.J_colind,
- d.efc.J,
+ dj.efc.J_rownnz,
+ dj.efc.J_rowadr,
+ dj.efc.J_colind,
+ dj.efc.J,
d.efc.D,
d.efc.state,
- ctx.changed_efc_ids,
- ctx.changed_efc_count,
+ dj.dof_cdof,
+ ctx.quad_changed_ids,
+ ctx.quad_changed_count,
slots,
],
outputs=[ctx.h],
@@ -3165,8 +3318,8 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
d.efc.J,
d.efc.D,
d.efc.state,
- ctx.changed_efc_ids,
- ctx.changed_efc_count,
+ ctx.quad_changed_ids,
+ ctx.quad_changed_count,
],
outputs=[ctx.h],
)
@@ -3250,72 +3403,6 @@ def _solve_beta_accumulate(
wp.atomic_add(ctx_beta_den_out, worldid, den)
-@cache_kernel
-def _solve_zero_search_dot(stable_fast: bool):
- STABLE_FAST = stable_fast
-
- @wp.kernel(module="unique", enable_backward=False)
- def kernel(
- # In:
- changed_count_in: wp.array[int],
- ctx_done_in: wp.array[bool],
- # Out:
- ctx_search_dot_out: wp.array[float],
- ):
- worldid = wp.tid()
-
- if ctx_done_in[worldid]:
- return
-
- # Fast path: search stays on the same ray; keep search_dot consistent with it.
- if wp.static(STABLE_FAST):
- if changed_count_in[worldid] == 0:
- return
-
- ctx_search_dot_out[worldid] = 0.0
-
- return kernel
-
-
-@cache_kernel
-def _solve_search_update(stable_fast: bool):
- STABLE_FAST = stable_fast
-
- @wp.kernel(module="unique", enable_backward=False)
- def kernel(
- # Model:
- opt_solver: int,
- # In:
- changed_count_in: wp.array[int],
- ctx_Mgrad_in: wp.array2d[float],
- ctx_search_in: wp.array2d[float],
- ctx_beta_in: wp.array[float],
- ctx_done_in: wp.array[bool],
- # Out:
- ctx_search_out: wp.array2d[float],
- ctx_search_dot_out: wp.array[float],
- ):
- worldid, dofid = wp.tid()
-
- if ctx_done_in[worldid]:
- return
-
- # Fast path: search stays on the stale ray; the linesearch absorbs its scale.
- if wp.static(STABLE_FAST):
- if changed_count_in[worldid] == 0:
- return
-
- search = -1.0 * ctx_Mgrad_in[worldid, dofid]
-
- if opt_solver == types.SolverType.CG:
- search += ctx_beta_in[worldid] * ctx_search_in[worldid, dofid]
-
- ctx_search_out[worldid, dofid] = search
- wp.atomic_add(ctx_search_dot_out, worldid, search * search)
-
- return kernel
-
-
@wp.kernel
def _solve_search_update_cg_tiled(
# Model:
@@ -3410,6 +3497,7 @@ def _solve_done(
stat_meaninertia: wp.array[float],
# In:
ctx_grad_dot_in: wp.array[float],
+ ctx_newton_decrement_in: wp.array[float],
ctx_improvement_in: wp.array[float],
ctx_done_in: wp.array[bool],
# Data out:
@@ -3429,7 +3517,8 @@ def _solve_done(
improvement = _rescale(nv, meaninertia, ctx_improvement_in[worldid])
gradient = _rescale(nv, meaninertia, wp.sqrt(ctx_grad_dot_in[worldid]))
- done = (improvement < tolerance) or (gradient < tolerance)
+ model_improvement = _rescale(nv, meaninertia, 0.5 * ctx_newton_decrement_in[worldid])
+ done = (improvement < tolerance) or (gradient < tolerance) or (model_improvement < tolerance)
if done or solver_niter_out[worldid] == opt_iterations:
# if the solver has converged or the maximum number of iterations has been reached then
# mark this world as done and remove it from the number of unconverged worlds
@@ -3437,14 +3526,29 @@ def _solve_done(
wp.atomic_add(nsolving_out, 0, -1)
+# The linesearch runs in ray units anchored at the last gradient rebuild, and
+# its bracketing arithmetic cannot resolve steps below ~eps times the gross
+# derivative-term magnitude over the curvature (see _linesearch_iterative_kernel).
+# A world whose tolerance sits below that floor could never terminate on the
+# stale ray; rebuilding re-anchors the arithmetic at the current (much smaller)
+# gradient scale. The 8 is an allowance for the reduction depth of the sums.
+_ALPHA_NOISE_EPS = 8.0 * 1.1920929e-07 # 8 * float32 eps
+
+
def _use_incremental(m: types.Model) -> bool:
"""Whether constraint state changes are tracked for incremental H updates."""
return m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC
-def _stable_fast(m: types.Model, compact: bool) -> bool:
- """Stable-state fast path: needs state-change tracking; compact scatters qfrc itself."""
- return _use_incremental(m) and not compact
+@wp.kernel(grid_stride=True)
+def _zero_change_counters(
+ # Out:
+ quad_changed_count_out: wp.array[int],
+ state_changed_count_out: wp.array[int],
+):
+ worldid = wp.tid()
+ quad_changed_count_out[worldid] = 0
+ state_changed_count_out[worldid] = 0
@event_scope
@@ -3462,20 +3566,23 @@ def _solver_iteration(
# tracking, and the additional JTCJ Hessian term depends on Jaref which
# changes every iteration.
incremental = _use_incremental(m)
- # Stable-state fast path: worlds with no state flips this iteration were
- # exactly quadratic over the step, so grad/Mgrad/search only changed by a
- # scalar along the same ray. Skip their qfrc/grad/solve/search updates and
- # track the scalar in ctx.grad_scale.
- stable_fast = _stable_fast(m, compact)
if incremental:
# Must complete before _update_constraint_efc which atomically increments.
- ctx.changed_efc_count.zero_()
+ wp.launch(
+ _zero_change_counters,
+ dim=d.nworld,
+ outputs=[ctx.quad_changed_count, ctx.state_changed_count],
+ )
- _update_constraint(m, d, ctx, track_changes=incremental, stable_fast=stable_fast)
+ # The tracking also enables the stable-state fast path: worlds with no state
+ # flips this iteration were exactly quadratic over the step, so grad/search
+ # only changed by a scalar along the same ray. Skip their qfrc/grad/
+ # solve/search updates and track the scalar in ctx.grad_scale.
+ _update_constraint(m, d, ctx, track_changes=incremental, stable_fast=incremental)
if incremental:
- _update_gradient_incremental(m, d, ctx, stable_fast)
+ _update_gradient_incremental(m, d, ctx, stable_fast=incremental)
else:
_update_gradient(m, d, ctx, compact=compact)
@@ -3523,16 +3630,6 @@ def _solver_iteration(
)
else:
- changed = ctx.changed_efc_count if stable_fast else d.nefc
- wp.launch(_solve_zero_search_dot(stable_fast), dim=d.nworld, inputs=[changed, ctx.done], outputs=[ctx.search_dot])
-
- wp.launch(
- _solve_search_update(stable_fast),
- dim=(d.nworld, m.nv),
- inputs=[m.opt.solver, changed, ctx.Mgrad, ctx.search, ctx.beta, ctx.done],
- outputs=[ctx.search, ctx.search_dot],
- )
-
wp.launch(
_solve_done,
dim=d.nworld,
@@ -3542,6 +3639,7 @@ def _solver_iteration(
m.opt.iterations,
m.stat.meaninertia,
ctx.grad_dot,
+ ctx.newton_decrement,
ctx.improvement,
ctx.done,
],
@@ -3576,18 +3674,39 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont
if threads_per_efc > 1:
ctx.Jaref.zero_()
+ sc = _sparse_compact(ctx)
+ dj = ctx.compact_d_full if sc else d
+ if sc:
+ dofs_per_thread = m.nv
+ threads_per_efc = 1
wp.launch(
- _solve_init_jaref_kernel(m.is_sparse, m.nv, dofs_per_thread),
+ _solve_init_jaref_kernel(sc or m.is_sparse, m.nv, dofs_per_thread, sc),
dim=(d.nworld, d.njmax, threads_per_efc),
- inputs=[d.nefc, d.qacc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.aref],
+ inputs=[d.nefc, d.qacc, dj.efc.J_rownnz, dj.efc.J_rowadr, dj.efc.J_colind, dj.efc.J, d.efc.aref, dj.dof_cdof],
outputs=[ctx.Jaref],
)
# Ma = M @ qacc
- support.mul_m(m, d, d.efc.Ma, d.qacc, skip=ctx.done)
+ _mul_m_compact_aware(m, d, ctx, d.efc.Ma, d.qacc, ctx.done)
_update_constraint(m, d, ctx)
+ # Build diagonal preconditioner (once per step).
+ if m.is_sparse and m.nflex > 0:
+ ctx.diag_precond = wp.empty(shape=(d.nworld, m.nv), dtype=float)
+ wp.launch(
+ _diag_precond_build,
+ dim=(d.nworld, m.nv),
+ inputs=[m.body_simple, m.dof_bodyid, m.M_rownnz, m.M_rowadr, d.M, ctx.done],
+ outputs=[ctx.diag_precond],
+ )
+ wp.launch(
+ _diag_precond_add_JTDJ,
+ dim=(d.nworld, d.njmax),
+ inputs=[m.body_simple, m.dof_bodyid, d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.D, ctx.done],
+ outputs=[ctx.diag_precond],
+ )
+
if grad:
_update_gradient(m, d, ctx, compact=compact)
@@ -3613,15 +3732,23 @@ def solve(m: types.Model, d: types.Data):
def _solve(m: types.Model, d: types.Data, ctx: SolverContext, compact: bool = False):
"""Finds forces that satisfy constraints."""
- if not (m.opt.disableflags & types.DisableBit.WARMSTART):
- wp.copy(d.qacc, d.qacc_warmstart)
- else:
- wp.copy(d.qacc, d.qacc_smooth)
+ warmstart = not (m.opt.disableflags & types.DisableBit.WARMSTART)
+ wp.launch(
+ _solve_init_dof(warmstart, m.is_sparse),
+ dim=(d.nworld, m.nv),
+ inputs=[d.nefc, d.qacc_warmstart, d.qacc_smooth],
+ outputs=[d.qacc, d.qfrc_constraint],
+ )
# context
init_context(m, d, ctx, grad=True, compact=compact)
- # search = -Mgrad
+ if _use_incremental(m):
+ # A new solve computes a new search direction: invalidate the mv/jv reuse
+ # left over from the previous solve.
+ ctx.search_unchanged.zero_()
+
+ # CG search = -Mgrad
if m.opt.solver == types.SolverType.CG:
wp.launch_tiled(
_solve_init_search_cg_tiled,
@@ -3631,14 +3758,6 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext, compact: bool = Fa
block_dim=m.block_dim.solve_init_search_cg,
)
- else:
- wp.launch(
- _solve_init_search,
- dim=(d.nworld, m.nv),
- inputs=[ctx.Mgrad],
- outputs=[ctx.search, ctx.search_dot],
- )
-
nsolving = wp.full(shape=(1,), value=d.nworld, dtype=int)
if m.opt.iterations != 0 and m.opt.graph_conditional:
# Note: the iteration kernel (indicated by while_body) is repeatedly launched
@@ -3656,9 +3775,10 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext, compact: bool = Fa
for _ in range(m.opt.iterations):
_solver_iteration(m, d, ctx, nsolving, compact=compact)
- # Recover the public qfrc_constraint: the fast path leaves it stale, and the
- # per-iteration zeroing wiped it for worlds that converged early.
- if _stable_fast(m, compact):
+ # Recover qfrc_constraint (the compacted buffer when run under solve_compact):
+ # the fast path leaves it stale, and the per-iteration zeroing wiped it for
+ # worlds that converged early.
+ if _use_incremental(m):
wp.launch(
_qfrc_constraint_from_grad,
dim=(d.nworld, m.nv),
@@ -3691,6 +3811,74 @@ def _init_compact_inertia(
M_c_out[worldid, i, j] = val
+@wp.kernel
+def _mul_m_sparse_compact(
+ # Model:
+ M_mulm_rowadr: wp.array[int],
+ M_mulm_col: wp.array[int],
+ M_mulm_madr: wp.array[int],
+ # Data in:
+ M_in: wp.array2d[float],
+ dof_cdof_in: wp.array2d[int],
+ cdof_dof_in: wp.array2d[int],
+ # In:
+ vec: wp.array2d[float],
+ skip: wp.array[bool],
+ # Out:
+ res: wp.array2d[float],
+):
+ """Compacted res = M @ vec via the full-coordinate sparse M (no gathered cM)."""
+ worldid, ci = wp.tid()
+
+ if skip[worldid]:
+ return
+
+ dof = cdof_dof_in[worldid, ci]
+ if dof < 0:
+ res[worldid, ci] = 0.0
+ return
+
+ acc = float(0.0)
+ start = M_mulm_rowadr[dof]
+ end = M_mulm_rowadr[dof + 1]
+ for k in range(start, end):
+ # tree-internal columns are awake with the row; guard like _gather_M_sparse
+ # in case M ever carries cross-tree entries (their gathered value is zero)
+ cj = dof_cdof_in[worldid, M_mulm_col[k]]
+ if cj >= 0:
+ acc += M_in[worldid, M_mulm_madr[k]] * vec[worldid, cj]
+ res[worldid, ci] = acc
+
+
+def _sparse_compact(ctx: SolverContext | InverseContext) -> bool:
+ """Whether this solve is compact over a sparse full model (full J structures exist)."""
+ return ctx.compact_d_full is not None and ctx.compact_m_full.is_sparse
+
+
+def _mul_m_compact_aware(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, res, vec, skip):
+ """M @ vec: full-coordinate sparse walk under compact, support.mul_m natively."""
+ dfull = ctx.compact_d_full
+ if dfull is not None:
+ mfull = ctx.compact_m_full
+ wp.launch(
+ _mul_m_sparse_compact,
+ dim=(d.nworld, m.nv),
+ inputs=[
+ mfull.M_mulm_rowadr,
+ mfull.M_mulm_col,
+ mfull.M_mulm_madr,
+ dfull.M,
+ dfull.dof_cdof,
+ dfull.cdof_dof,
+ vec,
+ skip,
+ ],
+ outputs=[res],
+ )
+ else:
+ support.mul_m(m, d, res, vec, skip=skip)
+
+
@wp.kernel
def _gather_M_sparse(
# Model:
@@ -3776,10 +3964,10 @@ def smooth_solve_compact(m: types.Model, d: types.Data):
outputs=[d.crhs],
)
wp.launch_tiled(
- _update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, d.nvmax_pad, False),
+ _cholesky_factorize_solve_blocked(types.TILE_SIZE_JTDAJ_DENSE, d.nvmax_pad),
dim=d.nworld,
- inputs=[wp.empty(0, dtype=bool), d.crhs, d.cM, d.cqLD],
- outputs=[d.cx],
+ inputs=[d.cM, d.crhs],
+ outputs=[d.cqLD, d.cx],
block_dim=m.block_dim.update_gradient_cholesky_blocked,
)
wp.launch(_scatter_solution, dim=(d.nworld, m.nv), inputs=[d.dof_cdof, d.cx], outputs=[d.qacc_smooth])
@@ -3881,7 +4069,8 @@ def solve_compact(m: types.Model, d: types.Data):
Gathers the active-DOF inertia, constraint Jacobian, and smooth/warmstart vectors
into nvmax_pad-sized dense workspaces, runs the stock dense Newton solver on a
shallow-replaced (m, d) at nvmax_pad, then scatters qacc/qfrc_constraint back.
- Inactive DOFs are frozen to 0. Reads the sparse Model inertia and constraint J.
+ Inactive DOFs are frozen to 0. On the incremental Newton path the solver
+ kernels read the sparse M and J directly through the compaction maps.
"""
_compact_gather(m, d)
@@ -3907,6 +4096,10 @@ def solve_compact(m: types.Model, d: types.Data):
)
sctx = _create_solver_context(m2, d2)
+ # compact kernels read the full-coordinate sparse structures (M, J) through
+ # the compaction maps instead of dense products on gathered blocks
+ sctx.compact_m_full = m
+ sctx.compact_d_full = d
_solve(m2, d2, sctx, compact=True)
_compact_scatter(m, d)
@@ -3915,30 +4108,22 @@ def solve_compact(m: types.Model, d: types.Data):
@event_scope
def _compact_gather(m: types.Model, d: types.Data):
nvp = d.nvmax_pad
- # gather compacted dense inertia (identity-padded tail)
- wp.launch(
- _init_compact_inertia,
- dim=(d.nworld, nvp, nvp),
- inputs=[d.ncdof],
- outputs=[d.cM],
- )
-
- wp.launch(
- _gather_M_sparse,
- dim=(d.nworld, m.nv),
- inputs=[m.M_rownnz, m.M_rowadr, m.M_colind, d.M, d.dof_cdof],
- outputs=[d.cM],
- )
- # gather compacted dense constraint Jacobian (active columns only)
- d.cJ.zero_()
- if m.is_sparse:
+ # gather compacted dense inertia and Jacobian only for dense models;
+ # sparse models read sparse M and J directly through compaction maps
+ if not m.is_sparse:
wp.launch(
- _gather_J_sparse,
- dim=(d.nworld, d.njmax),
- inputs=[d.nefc, d.dof_cdof, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J],
- outputs=[d.cJ],
+ _init_compact_inertia,
+ dim=(d.nworld, nvp, nvp),
+ inputs=[d.ncdof],
+ outputs=[d.cM],
)
- else:
+ wp.launch(
+ _gather_M_sparse,
+ dim=(d.nworld, m.nv),
+ inputs=[m.M_rownnz, m.M_rowadr, m.M_colind, d.M, d.dof_cdof],
+ outputs=[d.cM],
+ )
+ d.cJ.zero_()
wp.launch(
_gather_J_dense,
dim=(d.nworld, d.njmax),
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
index 4dc91e51..5177906f 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
@@ -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)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py
deleted file mode 100644
index 55a6a732..00000000
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py
+++ /dev/null
@@ -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 = """
-
-
-
-
-
-
-
-
-
-"""
-
-# 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 = """
-
-
-
-
-
-
-
-
- """
- 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"""
-
-
-
-
-
-
-
-
-
- """
- 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 = """
-
-
-
-
-
-
-
-
-
-
- """
- 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()
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
index 4b2b2803..91b43a8c 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
@@ -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
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
index 29db9ded..2182daf6 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
@@ -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",
]
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
index 388b9c48..a70827e3 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
@@ -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))
diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py
index 7f648694..8db84423 100644
--- a/mjx/mujoco/mjx/warp/bvh.py
+++ b/mjx/mujoco/mjx/warp/bvh.py
@@ -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,
diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py
index e5c95947..99d802c7 100644
--- a/mjx/mujoco/mjx/warp/collision_driver.py
+++ b/mjx/mujoco/mjx/warp/collision_driver.py
@@ -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,
)
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index a68351ff..6d0cb8e0 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -66,17 +66,17 @@ def _forward_shim(
actuator_actlimited: wp.array[bool],
actuator_actnum: wp.array[int],
actuator_actrange: wp.array2d[wp.vec2],
- actuator_biasprm: wp.array2d[mjwp_types.vec10f],
+ actuator_biasprm: wp.array2d[mjwp_types.vec10],
actuator_biastype: wp.array[int],
actuator_cranklength: wp.array2d[float],
actuator_ctrllimited: wp.array[bool],
actuator_ctrlrange: wp.array2d[wp.vec2],
actuator_delay: wp.array[float],
- actuator_dynprm: wp.array2d[mjwp_types.vec10f],
+ actuator_dynprm: wp.array2d[mjwp_types.vec10],
actuator_dyntype: wp.array[int],
actuator_forcelimited: wp.array[bool],
actuator_forcerange: wp.array2d[wp.vec2],
- actuator_gainprm: wp.array2d[mjwp_types.vec10f],
+ actuator_gainprm: wp.array2d[mjwp_types.vec10],
actuator_gaintype: wp.array[int],
actuator_gear: wp.array2d[wp.spatial_vector],
actuator_history: wp.array[wp.vec2i],
@@ -107,6 +107,7 @@ def _forward_shim(
body_pos: wp.array2d[wp.vec3],
body_quat: wp.array2d[wp.quat],
body_rootid: wp.array[int],
+ body_simple: wp.array[int],
body_subtreemass: wp.array2d[float],
body_tree: tuple[wp.array[int], ...],
body_treeid: wp.array[int],
@@ -150,6 +151,7 @@ def _forward_shim(
eq_ten_adr: wp.array[int],
eq_type: wp.array[int],
eq_wld_adr: wp.array[int],
+ flex_bend_interp_map: wp.array[wp.vec2i],
flex_bending: wp.array[float],
flex_bendingadr: wp.array[int],
flex_cell_map: wp.array[wp.vec4i],
@@ -176,6 +178,9 @@ def _forward_shim(
flex_evpairadr: wp.array[int],
flex_evpairflexid: wp.array[int],
flex_evpairnum: wp.array[int],
+ flex_face: wp.array2d[int],
+ flex_face_map: wp.array[wp.vec2i],
+ flex_faceadr: wp.array[int],
flex_friction: wp.array[wp.vec3],
flex_gap: wp.array[float],
flex_internal: wp.array[int],
@@ -264,7 +269,6 @@ def _forward_shim(
jnt_stiffness: wp.array2d[float],
jnt_stiffnesspoly: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
- jtcj_max_pairs: int,
light_bodyid: wp.array[int],
light_dir: wp.array2d[wp.vec3],
light_dir0: wp.array2d[wp.vec3],
@@ -306,9 +310,11 @@ def _forward_shim(
ncam: int,
neq: int,
nflex: int,
+ nflexbend_interp: int,
nflexedge: int,
nflexelem: int,
nflexevpair: int,
+ nflexface: int,
nflexintcell: int,
nflexnode: int,
nflexvert: int,
@@ -351,13 +357,7 @@ def _forward_shim(
qLD_all_updates: wp.array[wp.vec3i],
qLD_block_adr: wp.array[int],
qLD_block_total: int,
- qLD_dof_dense: wp.array[int],
- qLD_dof_simple: wp.array[int],
- qLD_has_dense: bool,
- qLD_has_simple: bool,
- qLD_has_sparse: bool,
qLD_level_offsets: wp.array[int],
- qLD_simple_dofs: wp.array[int],
qLD_updates: tuple[wp.array[wp.vec3i], ...],
qpos0: wp.array2d[float],
qpos_spring: wp.array2d[float],
@@ -512,6 +512,8 @@ def _forward_shim(
efc_islandid: wp.array2d[int],
energy: wp.array[wp.vec2],
eq_active: wp.array2d[bool],
+ face_quat: wp.array2d[wp.quat],
+ face_xpos: wp.array3d[wp.vec3],
flex_aabb_max: wp.array2d[wp.vec3],
flex_aabb_min: wp.array2d[wp.vec3],
flexedge_J: wp.array2d[float],
@@ -695,6 +697,7 @@ def _forward_shim(
_m.body_pos = body_pos
_m.body_quat = body_quat
_m.body_rootid = body_rootid
+ _m.body_simple = body_simple
_m.body_subtreemass = body_subtreemass
_m.body_tree = body_tree
_m.body_treeid = body_treeid
@@ -738,6 +741,7 @@ def _forward_shim(
_m.eq_ten_adr = eq_ten_adr
_m.eq_type = eq_type
_m.eq_wld_adr = eq_wld_adr
+ _m.flex_bend_interp_map = flex_bend_interp_map
_m.flex_bending = flex_bending
_m.flex_bendingadr = flex_bendingadr
_m.flex_cell_map = flex_cell_map
@@ -764,6 +768,9 @@ def _forward_shim(
_m.flex_evpairadr = flex_evpairadr
_m.flex_evpairflexid = flex_evpairflexid
_m.flex_evpairnum = flex_evpairnum
+ _m.flex_face = flex_face
+ _m.flex_face_map = flex_face_map
+ _m.flex_faceadr = flex_faceadr
_m.flex_friction = flex_friction
_m.flex_gap = flex_gap
_m.flex_internal = flex_internal
@@ -852,7 +859,6 @@ def _forward_shim(
_m.jnt_stiffness = jnt_stiffness
_m.jnt_stiffnesspoly = jnt_stiffnesspoly
_m.jnt_type = jnt_type
- _m.jtcj_max_pairs = jtcj_max_pairs
_m.light_bodyid = light_bodyid
_m.light_dir = light_dir
_m.light_dir0 = light_dir0
@@ -894,9 +900,11 @@ def _forward_shim(
_m.ncam = ncam
_m.neq = neq
_m.nflex = nflex
+ _m.nflexbend_interp = nflexbend_interp
_m.nflexedge = nflexedge
_m.nflexelem = nflexelem
_m.nflexevpair = nflexevpair
+ _m.nflexface = nflexface
_m.nflexintcell = nflexintcell
_m.nflexnode = nflexnode
_m.nflexvert = nflexvert
@@ -964,13 +972,7 @@ def _forward_shim(
_m.qLD_all_updates = qLD_all_updates
_m.qLD_block_adr = qLD_block_adr
_m.qLD_block_total = qLD_block_total
- _m.qLD_dof_dense = qLD_dof_dense
- _m.qLD_dof_simple = qLD_dof_simple
- _m.qLD_has_dense = qLD_has_dense
- _m.qLD_has_simple = qLD_has_simple
- _m.qLD_has_sparse = qLD_has_sparse
_m.qLD_level_offsets = qLD_level_offsets
- _m.qLD_simple_dofs = qLD_simple_dofs
_m.qLD_updates = qLD_updates
_m.qpos0 = qpos0
_m.qpos_spring = qpos_spring
@@ -1130,6 +1132,8 @@ def _forward_shim(
_d.efc_islandid = efc_islandid
_d.energy = energy
_d.eq_active = eq_active
+ _d.face_quat = face_quat
+ _d.face_xpos = face_xpos
_d.flex_aabb_max = flex_aabb_max
_d.flex_aabb_min = flex_aabb_min
_d.flexedge_J = flexedge_J
@@ -1245,6 +1249,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'cfrc_ext': d._impl.cfrc_ext.shape,
'cfrc_int': d._impl.cfrc_int.shape,
'cinert': d._impl.cinert.shape,
+ 'cqLD': d._impl.cqLD.shape,
'cqacc_smooth': d._impl.cqacc_smooth.shape,
'cqacc_warmstart': d._impl.cqacc_warmstart.shape,
'cqfrc_smooth': d._impl.cqfrc_smooth.shape,
@@ -1258,6 +1263,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'dof_islandid': d._impl.dof_islandid.shape,
'efc_islandid': d._impl.efc_islandid.shape,
'energy': d._impl.energy.shape,
+ 'face_quat': d._impl.face_quat.shape,
+ 'face_xpos': d._impl.face_xpos.shape,
'flex_aabb_max': d._impl.flex_aabb_max.shape,
'flex_aabb_min': d._impl.flex_aabb_min.shape,
'flexedge_J': d._impl.flexedge_J.shape,
@@ -1375,7 +1382,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_forward_shim,
- num_outputs=146,
+ num_outputs=149,
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([
@@ -1398,6 +1405,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'cfrc_ext',
'cfrc_int',
'cinert',
+ 'cqLD',
'cqacc_smooth',
'cqacc_warmstart',
'cqfrc_smooth',
@@ -1411,6 +1419,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'dof_islandid',
'efc_islandid',
'energy',
+ 'face_quat',
+ 'face_xpos',
'flex_aabb_max',
'flex_aabb_min',
'flexedge_J',
@@ -1527,6 +1537,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'efc__vel',
]),
stage_in_argnames=set([
+ 'M',
'act',
'act_dot',
'actuator_acc0',
@@ -1541,6 +1552,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'actuator_gear',
'actuator_length',
'actuator_lengthrange',
+ 'actuator_moment',
+ 'actuator_velocity',
+ 'body_awake',
+ 'body_awake_ind',
'body_gravcomp',
'body_inertia',
'body_invweight0',
@@ -1550,6 +1565,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'body_pos',
'body_quat',
'body_subtreemass',
+ 'cJ',
+ 'cM',
+ 'cMa',
+ 'cacc',
'cam_fovy',
'cam_intrinsic',
'cam_mat0',
@@ -1560,20 +1579,89 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'cam_xmat',
'cam_xpos',
'cdof',
+ 'cdof_dof',
'cdof_dot',
+ 'cdof_tri_col',
+ 'cdof_tri_row',
+ 'cfrc_ext',
+ 'cfrc_int',
+ 'cinert',
+ 'cls_tol',
+ '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',
+ 'cqLD',
+ 'cqacc',
+ 'cqacc_smooth',
+ 'cqacc_warmstart',
+ 'cqfrc_constraint',
+ 'cqfrc_smooth',
+ 'crb',
+ 'crhs',
+ 'ctol',
'ctrl',
'cvel',
+ 'cx',
'dof_armature',
+ 'dof_awake_ind',
+ 'dof_cdof',
'dof_damping',
'dof_dampingpoly',
'dof_frictionloss',
'dof_invweight0',
+ 'dof_island',
+ 'dof_islandid',
'dof_solimp',
'dof_solref',
+ 'efc__D',
+ 'efc__J',
+ 'efc__J_colind',
+ 'efc__J_rowadr',
+ 'efc__J_rownnz',
+ 'efc__Jqvel',
+ 'efc__Ma',
+ 'efc__aref',
+ 'efc__force',
+ 'efc__frictionloss',
+ 'efc__id',
+ 'efc__island',
+ 'efc__jtdaj_adr',
+ 'efc__jtdaj_nblock',
+ 'efc__jtdaj_nrow',
+ 'efc__margin',
+ 'efc__pos',
+ 'efc__state',
+ 'efc__type',
+ 'efc__vel',
+ 'efc_islandid',
+ 'energy',
'eq_active',
'eq_data',
'eq_solimp',
'eq_solref',
+ 'face_quat',
+ 'face_xpos',
+ 'flex_aabb_max',
+ 'flex_aabb_min',
+ 'flexedge_J',
+ 'flexedge_length',
+ 'flexedge_velocity',
+ 'flexnode_xpos',
+ 'flexvert_xpos',
'geom_aabb',
'geom_friction',
'geom_gap',
@@ -1591,6 +1679,13 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'geom_xpos',
'hfield_data',
'history',
+ 'island_dofadr',
+ 'island_idofadr',
+ 'island_iefcadr',
+ 'island_ne',
+ 'island_nefc',
+ 'island_nf',
+ 'island_nv',
'jnt_actfrcrange',
'jnt_axis',
'jnt_margin',
@@ -1605,23 +1700,49 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'light_pos',
'light_pos0',
'light_poscom0',
+ 'light_xdir',
+ 'light_xpos',
+ 'map_dof2idof',
+ 'map_efc2iefc',
+ 'map_idof2dof',
+ 'map_iefc2efc',
'mat_rgba',
'mocap_pos',
'mocap_quat',
+ 'moment_colind',
+ 'moment_rowadr',
+ 'moment_rownnz',
+ 'nacon',
+ 'nbody_awake',
+ 'ncdof',
+ 'ncollision',
+ 'ne',
+ 'nefc',
+ 'nf',
+ 'nidof',
+ 'nisland',
+ 'nl',
+ 'ntree_awake',
+ 'nv_awake',
+ 'opt__ccd_tolerance',
'opt__density',
'opt__gravity',
+ 'opt__impratio_invsqrt',
'opt__ls_tolerance',
'opt__magnetic',
'opt__timestep',
'opt__tolerance',
'opt__viscosity',
'opt__wind',
+ 'overflow',
'pair_friction',
'pair_gap',
'pair_margin',
'pair_solimp',
'pair_solref',
'pair_solreffriction',
+ 'qLD',
+ 'qLDiagInv',
'qacc',
'qacc_smooth',
'qacc_warmstart',
@@ -1629,10 +1750,12 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'qfrc_applied',
'qfrc_bias',
'qfrc_constraint',
+ 'qfrc_damper',
'qfrc_fluid',
'qfrc_gravcomp',
'qfrc_passive',
'qfrc_smooth',
+ 'qfrc_spring',
'qpos',
'qpos0',
'qpos_spring',
@@ -1642,8 +1765,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'site_quat',
'site_xmat',
'site_xpos',
+ 'solver_niter',
+ 'subtree_angmom',
'subtree_com',
+ 'subtree_linvel',
+ 'ten_J',
'ten_length',
+ 'ten_velocity',
+ 'ten_wrapadr',
+ 'ten_wrapnum',
'tendon_actfrcrange',
'tendon_armature',
'tendon_damping',
@@ -1661,6 +1791,11 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'tendon_stiffness',
'tendon_stiffnesspoly',
'time',
+ 'tree_asleep',
+ 'tree_awake',
+ 'tree_island',
+ 'wrap_obj',
+ 'wrap_xpos',
'xanchor',
'xaxis',
'xfrc_applied',
@@ -1671,32 +1806,148 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'xquat',
]),
stage_out_argnames=set([
+ 'M',
'act_dot',
'actuator_force',
'actuator_length',
+ 'actuator_moment',
+ 'actuator_velocity',
+ 'body_awake',
+ 'body_awake_ind',
+ 'cJ',
+ 'cM',
+ 'cacc',
'cam_xmat',
'cam_xpos',
'cdof',
+ 'cdof_dof',
'cdof_dot',
+ 'cfrc_ext',
+ 'cfrc_int',
+ 'cinert',
+ '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',
+ 'cqLD',
+ 'cqacc_smooth',
+ 'cqacc_warmstart',
+ 'cqfrc_smooth',
+ 'crb',
+ 'crhs',
'cvel',
+ 'cx',
+ 'dof_awake_ind',
+ 'dof_cdof',
+ 'dof_island',
+ 'dof_islandid',
+ 'efc__D',
+ 'efc__J',
+ 'efc__J_colind',
+ 'efc__J_rowadr',
+ 'efc__J_rownnz',
+ 'efc__Jqvel',
+ 'efc__Ma',
+ 'efc__aref',
+ 'efc__force',
+ 'efc__frictionloss',
+ 'efc__id',
+ 'efc__island',
+ 'efc__jtdaj_adr',
+ 'efc__jtdaj_nblock',
+ 'efc__jtdaj_nrow',
+ 'efc__margin',
+ 'efc__pos',
+ 'efc__state',
+ 'efc__type',
+ 'efc__vel',
+ 'efc_islandid',
+ 'energy',
+ 'face_quat',
+ 'face_xpos',
+ 'flex_aabb_max',
+ 'flex_aabb_min',
+ 'flexedge_J',
+ 'flexedge_length',
+ 'flexedge_velocity',
+ 'flexnode_xpos',
+ 'flexvert_xpos',
'geom_xmat',
'geom_xpos',
'history',
+ 'island_dofadr',
+ 'island_idofadr',
+ 'island_iefcadr',
+ 'island_ne',
+ 'island_nefc',
+ 'island_nf',
+ 'island_nv',
+ 'light_xdir',
+ 'light_xpos',
+ 'map_dof2idof',
+ 'map_efc2iefc',
+ 'map_idof2dof',
+ 'map_iefc2efc',
+ 'moment_colind',
+ 'moment_rowadr',
+ 'moment_rownnz',
+ 'nacon',
+ 'nbody_awake',
+ 'ncdof',
+ 'ncollision',
+ 'ne',
+ 'nefc',
+ 'nf',
+ 'nidof',
+ 'nisland',
+ 'nl',
+ 'ntree_awake',
+ 'nv_awake',
+ 'overflow',
+ 'qLD',
+ 'qLDiagInv',
'qacc',
'qacc_smooth',
'qfrc_actuator',
'qfrc_bias',
'qfrc_constraint',
+ 'qfrc_damper',
'qfrc_fluid',
'qfrc_gravcomp',
'qfrc_passive',
'qfrc_smooth',
+ 'qfrc_spring',
'qvel',
'sensordata',
'site_xmat',
'site_xpos',
+ 'solver_niter',
+ 'subtree_angmom',
'subtree_com',
+ 'subtree_linvel',
+ 'ten_J',
'ten_length',
+ 'ten_velocity',
+ 'ten_wrapadr',
+ 'ten_wrapnum',
+ 'tree_asleep',
+ 'tree_awake',
+ 'tree_island',
+ 'wrap_obj',
+ 'wrap_xpos',
'xanchor',
'xaxis',
'ximat',
@@ -1766,6 +2017,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.body_pos,
m.body_quat,
m.body_rootid,
+ m.body_simple,
m.body_subtreemass,
m._impl.body_tree,
m.body_treeid,
@@ -1809,6 +2061,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.eq_ten_adr,
m.eq_type,
m._impl.eq_wld_adr,
+ m._impl.flex_bend_interp_map,
m._impl.flex_bending,
m._impl.flex_bendingadr,
m._impl.flex_cell_map,
@@ -1835,6 +2088,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.flex_evpairadr,
m._impl.flex_evpairflexid,
m._impl.flex_evpairnum,
+ m._impl.flex_face,
+ m._impl.flex_face_map,
+ m._impl.flex_faceadr,
m._impl.flex_friction,
m._impl.flex_gap,
m._impl.flex_internal,
@@ -1923,7 +2179,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.jnt_stiffness,
m.jnt_stiffnesspoly,
m.jnt_type,
- m._impl.jtcj_max_pairs,
m._impl.light_bodyid,
m.light_dir,
m.light_dir0,
@@ -1965,9 +2220,11 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.ncam,
m.neq,
m.nflex,
+ m._impl.nflexbend_interp,
m._impl.nflexedge,
m._impl.nflexelem,
m._impl.nflexevpair,
+ m._impl.nflexface,
m._impl.nflexintcell,
m._impl.nflexnode,
m._impl.nflexvert,
@@ -2010,13 +2267,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m._impl.qLD_all_updates,
m._impl.qLD_block_adr,
m._impl.qLD_block_total,
- m._impl.qLD_dof_dense,
- m._impl.qLD_dof_simple,
- m._impl.qLD_has_dense,
- m._impl.qLD_has_simple,
- m._impl.qLD_has_sparse,
m._impl.qLD_level_offsets,
- m._impl.qLD_simple_dofs,
m._impl.qLD_updates,
m.qpos0,
m.qpos_spring,
@@ -2170,6 +2421,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d._impl.efc_islandid,
d._impl.energy,
d.eq_active,
+ d._impl.face_quat,
+ d._impl.face_xpos,
d._impl.flex_aabb_max,
d._impl.flex_aabb_min,
d._impl.flexedge_J,
@@ -2312,133 +2565,136 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'_impl.cfrc_ext': out[16],
'_impl.cfrc_int': out[17],
'_impl.cinert': out[18],
- '_impl.cqacc_smooth': out[19],
- '_impl.cqacc_warmstart': out[20],
- '_impl.cqfrc_smooth': out[21],
- '_impl.crb': out[22],
- '_impl.crhs': out[23],
- 'cvel': out[24],
- '_impl.cx': out[25],
- '_impl.dof_awake_ind': out[26],
- '_impl.dof_cdof': out[27],
- '_impl.dof_island': out[28],
- '_impl.dof_islandid': out[29],
- '_impl.efc_islandid': out[30],
- '_impl.energy': out[31],
- '_impl.flex_aabb_max': out[32],
- '_impl.flex_aabb_min': out[33],
- '_impl.flexedge_J': out[34],
- '_impl.flexedge_length': out[35],
- '_impl.flexedge_velocity': out[36],
- '_impl.flexnode_xpos': out[37],
- '_impl.flexvert_xpos': out[38],
- 'geom_xmat': out[39],
- 'geom_xpos': out[40],
- 'history': out[41],
- '_impl.island_dofadr': out[42],
- '_impl.island_idofadr': out[43],
- '_impl.island_iefcadr': out[44],
- '_impl.island_ne': out[45],
- '_impl.island_nefc': out[46],
- '_impl.island_nf': out[47],
- '_impl.island_nv': out[48],
- '_impl.light_xdir': out[49],
- '_impl.light_xpos': out[50],
- '_impl.map_dof2idof': out[51],
- '_impl.map_efc2iefc': out[52],
- '_impl.map_idof2dof': out[53],
- '_impl.map_iefc2efc': out[54],
- '_impl.moment_colind': out[55],
- '_impl.moment_rowadr': out[56],
- '_impl.moment_rownnz': out[57],
- '_impl.nacon': out[58],
- '_impl.nbody_awake': out[59],
- '_impl.ncdof': out[60],
- '_impl.ncollision': out[61],
- '_impl.ne': out[62],
- '_impl.nefc': out[63],
- '_impl.nf': out[64],
- '_impl.nidof': out[65],
- '_impl.nisland': out[66],
- '_impl.nl': out[67],
- '_impl.ntree_awake': out[68],
- '_impl.nv_awake': out[69],
- '_impl.overflow': out[70],
- '_impl.qLD': out[71],
- '_impl.qLDiagInv': out[72],
- 'qacc': out[73],
- 'qacc_smooth': out[74],
- 'qfrc_actuator': out[75],
- 'qfrc_bias': out[76],
- 'qfrc_constraint': out[77],
- '_impl.qfrc_damper': out[78],
- 'qfrc_fluid': out[79],
- 'qfrc_gravcomp': out[80],
- 'qfrc_passive': out[81],
- 'qfrc_smooth': out[82],
- '_impl.qfrc_spring': out[83],
- 'qvel': out[84],
- 'sensordata': out[85],
- 'site_xmat': out[86],
- 'site_xpos': out[87],
- '_impl.solver_niter': out[88],
- '_impl.subtree_angmom': out[89],
- 'subtree_com': out[90],
- '_impl.subtree_linvel': out[91],
- '_impl.ten_J': out[92],
- 'ten_length': out[93],
- '_impl.ten_velocity': out[94],
- '_impl.ten_wrapadr': out[95],
- '_impl.ten_wrapnum': out[96],
- '_impl.tree_asleep': out[97],
- '_impl.tree_awake': out[98],
- '_impl.tree_island': out[99],
- '_impl.wrap_obj': out[100],
- '_impl.wrap_xpos': out[101],
- 'xanchor': out[102],
- 'xaxis': out[103],
- 'ximat': out[104],
- 'xipos': out[105],
- 'xmat': out[106],
- 'xpos': out[107],
- 'xquat': out[108],
- '_impl.contact__dim': out[109],
- '_impl.contact__dist': out[110],
- '_impl.contact__efc_address': out[111],
- '_impl.contact__elem': out[112],
- '_impl.contact__flex': out[113],
- '_impl.contact__frame': out[114],
- '_impl.contact__friction': out[115],
- '_impl.contact__geom': out[116],
- '_impl.contact__geomcollisionid': out[117],
- '_impl.contact__includemargin': out[118],
- '_impl.contact__pos': out[119],
- '_impl.contact__solimp': out[120],
- '_impl.contact__solref': out[121],
- '_impl.contact__solreffriction': out[122],
- '_impl.contact__type': out[123],
- '_impl.contact__vert': out[124],
- '_impl.contact__worldid': out[125],
- '_impl.efc__D': out[126],
- '_impl.efc__J': out[127],
- '_impl.efc__J_colind': out[128],
- '_impl.efc__J_rowadr': out[129],
- '_impl.efc__J_rownnz': out[130],
- '_impl.efc__Jqvel': out[131],
- '_impl.efc__Ma': out[132],
- '_impl.efc__aref': out[133],
- '_impl.efc__force': out[134],
- '_impl.efc__frictionloss': out[135],
- '_impl.efc__id': out[136],
- '_impl.efc__island': out[137],
- '_impl.efc__jtdaj_adr': out[138],
- '_impl.efc__jtdaj_nblock': out[139],
- '_impl.efc__jtdaj_nrow': out[140],
- '_impl.efc__margin': out[141],
- '_impl.efc__pos': out[142],
- '_impl.efc__state': out[143],
- '_impl.efc__type': out[144],
- '_impl.efc__vel': out[145],
+ '_impl.cqLD': out[19],
+ '_impl.cqacc_smooth': out[20],
+ '_impl.cqacc_warmstart': out[21],
+ '_impl.cqfrc_smooth': out[22],
+ '_impl.crb': out[23],
+ '_impl.crhs': out[24],
+ 'cvel': out[25],
+ '_impl.cx': out[26],
+ '_impl.dof_awake_ind': out[27],
+ '_impl.dof_cdof': out[28],
+ '_impl.dof_island': out[29],
+ '_impl.dof_islandid': out[30],
+ '_impl.efc_islandid': out[31],
+ '_impl.energy': out[32],
+ '_impl.face_quat': out[33],
+ '_impl.face_xpos': out[34],
+ '_impl.flex_aabb_max': out[35],
+ '_impl.flex_aabb_min': out[36],
+ '_impl.flexedge_J': out[37],
+ '_impl.flexedge_length': out[38],
+ '_impl.flexedge_velocity': out[39],
+ '_impl.flexnode_xpos': out[40],
+ '_impl.flexvert_xpos': out[41],
+ 'geom_xmat': out[42],
+ 'geom_xpos': out[43],
+ 'history': out[44],
+ '_impl.island_dofadr': out[45],
+ '_impl.island_idofadr': out[46],
+ '_impl.island_iefcadr': out[47],
+ '_impl.island_ne': out[48],
+ '_impl.island_nefc': out[49],
+ '_impl.island_nf': out[50],
+ '_impl.island_nv': out[51],
+ '_impl.light_xdir': out[52],
+ '_impl.light_xpos': out[53],
+ '_impl.map_dof2idof': out[54],
+ '_impl.map_efc2iefc': out[55],
+ '_impl.map_idof2dof': out[56],
+ '_impl.map_iefc2efc': out[57],
+ '_impl.moment_colind': out[58],
+ '_impl.moment_rowadr': out[59],
+ '_impl.moment_rownnz': out[60],
+ '_impl.nacon': out[61],
+ '_impl.nbody_awake': out[62],
+ '_impl.ncdof': out[63],
+ '_impl.ncollision': out[64],
+ '_impl.ne': out[65],
+ '_impl.nefc': out[66],
+ '_impl.nf': out[67],
+ '_impl.nidof': out[68],
+ '_impl.nisland': out[69],
+ '_impl.nl': out[70],
+ '_impl.ntree_awake': out[71],
+ '_impl.nv_awake': out[72],
+ '_impl.overflow': out[73],
+ '_impl.qLD': out[74],
+ '_impl.qLDiagInv': out[75],
+ 'qacc': out[76],
+ 'qacc_smooth': out[77],
+ 'qfrc_actuator': out[78],
+ 'qfrc_bias': out[79],
+ 'qfrc_constraint': out[80],
+ '_impl.qfrc_damper': out[81],
+ 'qfrc_fluid': out[82],
+ 'qfrc_gravcomp': out[83],
+ 'qfrc_passive': out[84],
+ 'qfrc_smooth': out[85],
+ '_impl.qfrc_spring': out[86],
+ 'qvel': out[87],
+ 'sensordata': out[88],
+ 'site_xmat': out[89],
+ 'site_xpos': out[90],
+ '_impl.solver_niter': out[91],
+ '_impl.subtree_angmom': out[92],
+ 'subtree_com': out[93],
+ '_impl.subtree_linvel': out[94],
+ '_impl.ten_J': out[95],
+ 'ten_length': out[96],
+ '_impl.ten_velocity': out[97],
+ '_impl.ten_wrapadr': out[98],
+ '_impl.ten_wrapnum': out[99],
+ '_impl.tree_asleep': out[100],
+ '_impl.tree_awake': out[101],
+ '_impl.tree_island': out[102],
+ '_impl.wrap_obj': out[103],
+ '_impl.wrap_xpos': out[104],
+ 'xanchor': out[105],
+ 'xaxis': out[106],
+ 'ximat': out[107],
+ 'xipos': out[108],
+ 'xmat': out[109],
+ 'xpos': out[110],
+ 'xquat': out[111],
+ '_impl.contact__dim': out[112],
+ '_impl.contact__dist': out[113],
+ '_impl.contact__efc_address': out[114],
+ '_impl.contact__elem': out[115],
+ '_impl.contact__flex': out[116],
+ '_impl.contact__frame': out[117],
+ '_impl.contact__friction': out[118],
+ '_impl.contact__geom': out[119],
+ '_impl.contact__geomcollisionid': out[120],
+ '_impl.contact__includemargin': out[121],
+ '_impl.contact__pos': out[122],
+ '_impl.contact__solimp': out[123],
+ '_impl.contact__solref': out[124],
+ '_impl.contact__solreffriction': out[125],
+ '_impl.contact__type': out[126],
+ '_impl.contact__vert': out[127],
+ '_impl.contact__worldid': out[128],
+ '_impl.efc__D': out[129],
+ '_impl.efc__J': out[130],
+ '_impl.efc__J_colind': out[131],
+ '_impl.efc__J_rowadr': out[132],
+ '_impl.efc__J_rownnz': out[133],
+ '_impl.efc__Jqvel': out[134],
+ '_impl.efc__Ma': out[135],
+ '_impl.efc__aref': out[136],
+ '_impl.efc__force': out[137],
+ '_impl.efc__frictionloss': out[138],
+ '_impl.efc__id': out[139],
+ '_impl.efc__island': out[140],
+ '_impl.efc__jtdaj_adr': out[141],
+ '_impl.efc__jtdaj_nblock': out[142],
+ '_impl.efc__jtdaj_nrow': out[143],
+ '_impl.efc__margin': out[144],
+ '_impl.efc__pos': out[145],
+ '_impl.efc__state': out[146],
+ '_impl.efc__type': out[147],
+ '_impl.efc__vel': out[148],
})
return d
@@ -2481,17 +2737,17 @@ def _step_shim(
actuator_actlimited: wp.array[bool],
actuator_actnum: wp.array[int],
actuator_actrange: wp.array2d[wp.vec2],
- actuator_biasprm: wp.array2d[mjwp_types.vec10f],
+ actuator_biasprm: wp.array2d[mjwp_types.vec10],
actuator_biastype: wp.array[int],
actuator_cranklength: wp.array2d[float],
actuator_ctrllimited: wp.array[bool],
actuator_ctrlrange: wp.array2d[wp.vec2],
actuator_delay: wp.array[float],
- actuator_dynprm: wp.array2d[mjwp_types.vec10f],
+ actuator_dynprm: wp.array2d[mjwp_types.vec10],
actuator_dyntype: wp.array[int],
actuator_forcelimited: wp.array[bool],
actuator_forcerange: wp.array2d[wp.vec2],
- actuator_gainprm: wp.array2d[mjwp_types.vec10f],
+ actuator_gainprm: wp.array2d[mjwp_types.vec10],
actuator_gaintype: wp.array[int],
actuator_gear: wp.array2d[wp.spatial_vector],
actuator_history: wp.array[wp.vec2i],
@@ -2524,6 +2780,7 @@ def _step_shim(
body_pos: wp.array2d[wp.vec3],
body_quat: wp.array2d[wp.quat],
body_rootid: wp.array[int],
+ body_simple: wp.array[int],
body_subtreemass: wp.array2d[float],
body_tree: tuple[wp.array[int], ...],
body_treeid: wp.array[int],
@@ -2567,6 +2824,7 @@ def _step_shim(
eq_ten_adr: wp.array[int],
eq_type: wp.array[int],
eq_wld_adr: wp.array[int],
+ flex_bend_interp_map: wp.array[wp.vec2i],
flex_bending: wp.array[float],
flex_bendingadr: wp.array[int],
flex_cell_map: wp.array[wp.vec4i],
@@ -2593,6 +2851,9 @@ def _step_shim(
flex_evpairadr: wp.array[int],
flex_evpairflexid: wp.array[int],
flex_evpairnum: wp.array[int],
+ flex_face: wp.array2d[int],
+ flex_face_map: wp.array[wp.vec2i],
+ flex_faceadr: wp.array[int],
flex_friction: wp.array[wp.vec3],
flex_gap: wp.array[float],
flex_internal: wp.array[int],
@@ -2681,7 +2942,6 @@ def _step_shim(
jnt_stiffness: wp.array2d[float],
jnt_stiffnesspoly: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
- jtcj_max_pairs: int,
light_bodyid: wp.array[int],
light_dir: wp.array2d[wp.vec3],
light_dir0: wp.array2d[wp.vec3],
@@ -2726,9 +2986,11 @@ def _step_shim(
ncam: int,
neq: int,
nflex: int,
+ nflexbend_interp: int,
nflexedge: int,
nflexelem: int,
nflexevpair: int,
+ nflexface: int,
nflexintcell: int,
nflexnode: int,
nflexvert: int,
@@ -2773,13 +3035,7 @@ def _step_shim(
qLD_all_updates: wp.array[wp.vec3i],
qLD_block_adr: wp.array[int],
qLD_block_total: int,
- qLD_dof_dense: wp.array[int],
- qLD_dof_simple: wp.array[int],
- qLD_has_dense: bool,
- qLD_has_simple: bool,
- qLD_has_sparse: bool,
qLD_level_offsets: wp.array[int],
- qLD_simple_dofs: wp.array[int],
qLD_updates: tuple[wp.array[wp.vec3i], ...],
qpos0: wp.array2d[float],
qpos_spring: wp.array2d[float],
@@ -2936,6 +3192,8 @@ def _step_shim(
efc_islandid: wp.array2d[int],
energy: wp.array[wp.vec2],
eq_active: wp.array2d[bool],
+ face_quat: wp.array2d[wp.quat],
+ face_xpos: wp.array3d[wp.vec3],
flex_aabb_max: wp.array2d[wp.vec3],
flex_aabb_min: wp.array2d[wp.vec3],
flexedge_J: wp.array2d[float],
@@ -3128,6 +3386,7 @@ def _step_shim(
_m.body_pos = body_pos
_m.body_quat = body_quat
_m.body_rootid = body_rootid
+ _m.body_simple = body_simple
_m.body_subtreemass = body_subtreemass
_m.body_tree = body_tree
_m.body_treeid = body_treeid
@@ -3171,6 +3430,7 @@ def _step_shim(
_m.eq_ten_adr = eq_ten_adr
_m.eq_type = eq_type
_m.eq_wld_adr = eq_wld_adr
+ _m.flex_bend_interp_map = flex_bend_interp_map
_m.flex_bending = flex_bending
_m.flex_bendingadr = flex_bendingadr
_m.flex_cell_map = flex_cell_map
@@ -3197,6 +3457,9 @@ def _step_shim(
_m.flex_evpairadr = flex_evpairadr
_m.flex_evpairflexid = flex_evpairflexid
_m.flex_evpairnum = flex_evpairnum
+ _m.flex_face = flex_face
+ _m.flex_face_map = flex_face_map
+ _m.flex_faceadr = flex_faceadr
_m.flex_friction = flex_friction
_m.flex_gap = flex_gap
_m.flex_internal = flex_internal
@@ -3285,7 +3548,6 @@ def _step_shim(
_m.jnt_stiffness = jnt_stiffness
_m.jnt_stiffnesspoly = jnt_stiffnesspoly
_m.jnt_type = jnt_type
- _m.jtcj_max_pairs = jtcj_max_pairs
_m.light_bodyid = light_bodyid
_m.light_dir = light_dir
_m.light_dir0 = light_dir0
@@ -3330,9 +3592,11 @@ def _step_shim(
_m.ncam = ncam
_m.neq = neq
_m.nflex = nflex
+ _m.nflexbend_interp = nflexbend_interp
_m.nflexedge = nflexedge
_m.nflexelem = nflexelem
_m.nflexevpair = nflexevpair
+ _m.nflexface = nflexface
_m.nflexintcell = nflexintcell
_m.nflexnode = nflexnode
_m.nflexvert = nflexvert
@@ -3404,13 +3668,7 @@ def _step_shim(
_m.qLD_all_updates = qLD_all_updates
_m.qLD_block_adr = qLD_block_adr
_m.qLD_block_total = qLD_block_total
- _m.qLD_dof_dense = qLD_dof_dense
- _m.qLD_dof_simple = qLD_dof_simple
- _m.qLD_has_dense = qLD_has_dense
- _m.qLD_has_simple = qLD_has_simple
- _m.qLD_has_sparse = qLD_has_sparse
_m.qLD_level_offsets = qLD_level_offsets
- _m.qLD_simple_dofs = qLD_simple_dofs
_m.qLD_updates = qLD_updates
_m.qpos0 = qpos0
_m.qpos_spring = qpos_spring
@@ -3570,6 +3828,8 @@ def _step_shim(
_d.efc_islandid = efc_islandid
_d.energy = energy
_d.eq_active = eq_active
+ _d.face_quat = face_quat
+ _d.face_xpos = face_xpos
_d.flex_aabb_max = flex_aabb_max
_d.flex_aabb_min = flex_aabb_min
_d.flexedge_J = flexedge_J
@@ -3687,6 +3947,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'cfrc_ext': d._impl.cfrc_ext.shape,
'cfrc_int': d._impl.cfrc_int.shape,
'cinert': d._impl.cinert.shape,
+ 'cqLD': d._impl.cqLD.shape,
'cqacc_smooth': d._impl.cqacc_smooth.shape,
'cqacc_warmstart': d._impl.cqacc_warmstart.shape,
'cqfrc_smooth': d._impl.cqfrc_smooth.shape,
@@ -3700,6 +3961,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'dof_islandid': d._impl.dof_islandid.shape,
'efc_islandid': d._impl.efc_islandid.shape,
'energy': d._impl.energy.shape,
+ 'face_quat': d._impl.face_quat.shape,
+ 'face_xpos': d._impl.face_xpos.shape,
'flex_aabb_max': d._impl.flex_aabb_max.shape,
'flex_aabb_min': d._impl.flex_aabb_min.shape,
'flexedge_J': d._impl.flexedge_J.shape,
@@ -3821,7 +4084,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_step_shim,
- num_outputs=151,
+ num_outputs=154,
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([
@@ -3845,6 +4108,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'cfrc_ext',
'cfrc_int',
'cinert',
+ 'cqLD',
'cqacc_smooth',
'cqacc_warmstart',
'cqfrc_smooth',
@@ -3858,6 +4122,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'dof_islandid',
'efc_islandid',
'energy',
+ 'face_quat',
+ 'face_xpos',
'flex_aabb_max',
'flex_aabb_min',
'flexedge_J',
@@ -3978,6 +4244,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'efc__vel',
]),
stage_in_argnames=set([
+ 'M',
'act',
'act_dot',
'actuator_acc0',
@@ -3992,6 +4259,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'actuator_gear',
'actuator_length',
'actuator_lengthrange',
+ 'actuator_moment',
+ 'actuator_velocity',
+ 'body_awake',
+ 'body_awake_ind',
'body_gravcomp',
'body_inertia',
'body_invweight0',
@@ -4001,6 +4272,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'body_pos',
'body_quat',
'body_subtreemass',
+ 'cJ',
+ 'cM',
+ 'cMa',
+ 'cacc',
'cam_fovy',
'cam_intrinsic',
'cam_mat0',
@@ -4011,20 +4286,89 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'cam_xmat',
'cam_xpos',
'cdof',
+ 'cdof_dof',
'cdof_dot',
+ 'cdof_tri_col',
+ 'cdof_tri_row',
+ 'cfrc_ext',
+ 'cfrc_int',
+ 'cinert',
+ 'cls_tol',
+ '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',
+ 'cqLD',
+ 'cqacc',
+ 'cqacc_smooth',
+ 'cqacc_warmstart',
+ 'cqfrc_constraint',
+ 'cqfrc_smooth',
+ 'crb',
+ 'crhs',
+ 'ctol',
'ctrl',
'cvel',
+ 'cx',
'dof_armature',
+ 'dof_awake_ind',
+ 'dof_cdof',
'dof_damping',
'dof_dampingpoly',
'dof_frictionloss',
'dof_invweight0',
+ 'dof_island',
+ 'dof_islandid',
'dof_solimp',
'dof_solref',
+ 'efc__D',
+ 'efc__J',
+ 'efc__J_colind',
+ 'efc__J_rowadr',
+ 'efc__J_rownnz',
+ 'efc__Jqvel',
+ 'efc__Ma',
+ 'efc__aref',
+ 'efc__force',
+ 'efc__frictionloss',
+ 'efc__id',
+ 'efc__island',
+ 'efc__jtdaj_adr',
+ 'efc__jtdaj_nblock',
+ 'efc__jtdaj_nrow',
+ 'efc__margin',
+ 'efc__pos',
+ 'efc__state',
+ 'efc__type',
+ 'efc__vel',
+ 'efc_islandid',
+ 'energy',
'eq_active',
'eq_data',
'eq_solimp',
'eq_solref',
+ 'face_quat',
+ 'face_xpos',
+ 'flex_aabb_max',
+ 'flex_aabb_min',
+ 'flexedge_J',
+ 'flexedge_length',
+ 'flexedge_velocity',
+ 'flexnode_xpos',
+ 'flexvert_xpos',
'geom_aabb',
'geom_friction',
'geom_gap',
@@ -4042,6 +4386,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'geom_xpos',
'hfield_data',
'history',
+ 'island_dofadr',
+ 'island_idofadr',
+ 'island_iefcadr',
+ 'island_ne',
+ 'island_nefc',
+ 'island_nf',
+ 'island_nv',
'jnt_actfrcrange',
'jnt_axis',
'jnt_margin',
@@ -4056,23 +4407,51 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'light_pos',
'light_pos0',
'light_poscom0',
+ 'light_xdir',
+ 'light_xpos',
+ 'map_dof2idof',
+ 'map_efc2iefc',
+ 'map_idof2dof',
+ 'map_iefc2efc',
'mat_rgba',
'mocap_pos',
'mocap_quat',
+ 'moment_colind',
+ 'moment_rowadr',
+ 'moment_rownnz',
+ 'nacon',
+ 'nbody_awake',
+ 'ncdof',
+ 'ncollision',
+ 'ne',
+ 'nefc',
+ 'nf',
+ 'nidof',
+ 'nisland',
+ 'nl',
+ 'ntree_awake',
+ 'nv_awake',
+ 'opt__ccd_tolerance',
'opt__density',
'opt__gravity',
+ 'opt__impratio_invsqrt',
'opt__ls_tolerance',
'opt__magnetic',
+ 'opt__sleep_tolerance',
'opt__timestep',
'opt__tolerance',
'opt__viscosity',
'opt__wind',
+ 'overflow',
'pair_friction',
'pair_gap',
'pair_margin',
'pair_solimp',
'pair_solref',
'pair_solreffriction',
+ 'qLD',
+ 'qLDiagInv',
+ 'qLU',
'qacc',
'qacc_smooth',
'qacc_warmstart',
@@ -4080,10 +4459,12 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'qfrc_applied',
'qfrc_bias',
'qfrc_constraint',
+ 'qfrc_damper',
'qfrc_fluid',
'qfrc_gravcomp',
'qfrc_passive',
'qfrc_smooth',
+ 'qfrc_spring',
'qpos',
'qpos0',
'qpos_spring',
@@ -4093,8 +4474,15 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'site_quat',
'site_xmat',
'site_xpos',
+ 'solver_niter',
+ 'subtree_angmom',
'subtree_com',
+ 'subtree_linvel',
+ 'ten_J',
'ten_length',
+ 'ten_velocity',
+ 'ten_wrapadr',
+ 'ten_wrapnum',
'tendon_actfrcrange',
'tendon_armature',
'tendon_damping',
@@ -4112,6 +4500,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'tendon_stiffness',
'tendon_stiffnesspoly',
'time',
+ 'tree_asleep',
+ 'tree_awake',
+ 'tree_island',
+ 'wrap_obj',
+ 'wrap_xpos',
'xanchor',
'xaxis',
'xfrc_applied',
@@ -4122,36 +4515,153 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'xquat',
]),
stage_out_argnames=set([
+ 'M',
'act',
'act_dot',
'actuator_force',
'actuator_length',
+ 'actuator_moment',
+ 'actuator_velocity',
+ 'body_awake',
+ 'body_awake_ind',
+ 'cJ',
+ 'cM',
+ 'cacc',
'cam_xmat',
'cam_xpos',
'cdof',
+ 'cdof_dof',
'cdof_dot',
+ 'cfrc_ext',
+ 'cfrc_int',
+ 'cinert',
+ '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',
+ 'cqLD',
+ 'cqacc_smooth',
+ 'cqacc_warmstart',
+ 'cqfrc_smooth',
+ 'crb',
+ 'crhs',
'cvel',
+ 'cx',
+ 'dof_awake_ind',
+ 'dof_cdof',
+ 'dof_island',
+ 'dof_islandid',
+ 'efc__D',
+ 'efc__J',
+ 'efc__J_colind',
+ 'efc__J_rowadr',
+ 'efc__J_rownnz',
+ 'efc__Jqvel',
+ 'efc__Ma',
+ 'efc__aref',
+ 'efc__force',
+ 'efc__frictionloss',
+ 'efc__id',
+ 'efc__island',
+ 'efc__jtdaj_adr',
+ 'efc__jtdaj_nblock',
+ 'efc__jtdaj_nrow',
+ 'efc__margin',
+ 'efc__pos',
+ 'efc__state',
+ 'efc__type',
+ 'efc__vel',
+ 'efc_islandid',
+ 'energy',
+ 'face_quat',
+ 'face_xpos',
+ 'flex_aabb_max',
+ 'flex_aabb_min',
+ 'flexedge_J',
+ 'flexedge_length',
+ 'flexedge_velocity',
+ 'flexnode_xpos',
+ 'flexvert_xpos',
'geom_xmat',
'geom_xpos',
'history',
+ 'island_dofadr',
+ 'island_idofadr',
+ 'island_iefcadr',
+ 'island_ne',
+ 'island_nefc',
+ 'island_nf',
+ 'island_nv',
+ 'light_xdir',
+ 'light_xpos',
+ 'map_dof2idof',
+ 'map_efc2iefc',
+ 'map_idof2dof',
+ 'map_iefc2efc',
+ 'moment_colind',
+ 'moment_rowadr',
+ 'moment_rownnz',
+ 'nacon',
+ 'nbody_awake',
+ 'ncdof',
+ 'ncollision',
+ 'ne',
+ 'nefc',
+ 'nf',
+ 'nidof',
+ 'nisland',
+ 'nl',
+ 'ntree_awake',
+ 'nv_awake',
+ 'overflow',
+ 'qLD',
+ 'qLDiagInv',
+ 'qLU',
'qacc',
'qacc_smooth',
'qacc_warmstart',
'qfrc_actuator',
'qfrc_bias',
'qfrc_constraint',
+ 'qfrc_damper',
'qfrc_fluid',
'qfrc_gravcomp',
'qfrc_passive',
'qfrc_smooth',
+ 'qfrc_spring',
'qpos',
'qvel',
'sensordata',
'site_xmat',
'site_xpos',
+ 'solver_niter',
+ 'subtree_angmom',
'subtree_com',
+ 'subtree_linvel',
+ 'ten_J',
'ten_length',
+ 'ten_velocity',
+ 'ten_wrapadr',
+ 'ten_wrapnum',
'time',
+ 'tree_asleep',
+ 'tree_awake',
+ 'tree_island',
+ 'wrap_obj',
+ 'wrap_xpos',
'xanchor',
'xaxis',
'ximat',
@@ -4229,6 +4739,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.body_pos,
m.body_quat,
m.body_rootid,
+ m.body_simple,
m.body_subtreemass,
m._impl.body_tree,
m.body_treeid,
@@ -4272,6 +4783,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.eq_ten_adr,
m.eq_type,
m._impl.eq_wld_adr,
+ m._impl.flex_bend_interp_map,
m._impl.flex_bending,
m._impl.flex_bendingadr,
m._impl.flex_cell_map,
@@ -4298,6 +4810,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.flex_evpairadr,
m._impl.flex_evpairflexid,
m._impl.flex_evpairnum,
+ m._impl.flex_face,
+ m._impl.flex_face_map,
+ m._impl.flex_faceadr,
m._impl.flex_friction,
m._impl.flex_gap,
m._impl.flex_internal,
@@ -4386,7 +4901,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.jnt_stiffness,
m.jnt_stiffnesspoly,
m.jnt_type,
- m._impl.jtcj_max_pairs,
m._impl.light_bodyid,
m.light_dir,
m.light_dir0,
@@ -4431,9 +4945,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.ncam,
m.neq,
m.nflex,
+ m._impl.nflexbend_interp,
m._impl.nflexedge,
m._impl.nflexelem,
m._impl.nflexevpair,
+ m._impl.nflexface,
m._impl.nflexintcell,
m._impl.nflexnode,
m._impl.nflexvert,
@@ -4478,13 +4994,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m._impl.qLD_all_updates,
m._impl.qLD_block_adr,
m._impl.qLD_block_total,
- m._impl.qLD_dof_dense,
- m._impl.qLD_dof_simple,
- m._impl.qLD_has_dense,
- m._impl.qLD_has_simple,
- m._impl.qLD_has_sparse,
m._impl.qLD_level_offsets,
- m._impl.qLD_simple_dofs,
m._impl.qLD_updates,
m.qpos0,
m.qpos_spring,
@@ -4640,6 +5150,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d._impl.efc_islandid,
d._impl.energy,
d.eq_active,
+ d._impl.face_quat,
+ d._impl.face_xpos,
d._impl.flex_aabb_max,
d._impl.flex_aabb_min,
d._impl.flexedge_J,
@@ -4784,137 +5296,140 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'_impl.cfrc_ext': out[17],
'_impl.cfrc_int': out[18],
'_impl.cinert': out[19],
- '_impl.cqacc_smooth': out[20],
- '_impl.cqacc_warmstart': out[21],
- '_impl.cqfrc_smooth': out[22],
- '_impl.crb': out[23],
- '_impl.crhs': out[24],
- 'cvel': out[25],
- '_impl.cx': out[26],
- '_impl.dof_awake_ind': out[27],
- '_impl.dof_cdof': out[28],
- '_impl.dof_island': out[29],
- '_impl.dof_islandid': out[30],
- '_impl.efc_islandid': out[31],
- '_impl.energy': out[32],
- '_impl.flex_aabb_max': out[33],
- '_impl.flex_aabb_min': out[34],
- '_impl.flexedge_J': out[35],
- '_impl.flexedge_length': out[36],
- '_impl.flexedge_velocity': out[37],
- '_impl.flexnode_xpos': out[38],
- '_impl.flexvert_xpos': out[39],
- 'geom_xmat': out[40],
- 'geom_xpos': out[41],
- 'history': out[42],
- '_impl.island_dofadr': out[43],
- '_impl.island_idofadr': out[44],
- '_impl.island_iefcadr': out[45],
- '_impl.island_ne': out[46],
- '_impl.island_nefc': out[47],
- '_impl.island_nf': out[48],
- '_impl.island_nv': out[49],
- '_impl.light_xdir': out[50],
- '_impl.light_xpos': out[51],
- '_impl.map_dof2idof': out[52],
- '_impl.map_efc2iefc': out[53],
- '_impl.map_idof2dof': out[54],
- '_impl.map_iefc2efc': out[55],
- '_impl.moment_colind': out[56],
- '_impl.moment_rowadr': out[57],
- '_impl.moment_rownnz': out[58],
- '_impl.nacon': out[59],
- '_impl.nbody_awake': out[60],
- '_impl.ncdof': out[61],
- '_impl.ncollision': out[62],
- '_impl.ne': out[63],
- '_impl.nefc': out[64],
- '_impl.nf': out[65],
- '_impl.nidof': out[66],
- '_impl.nisland': out[67],
- '_impl.nl': out[68],
- '_impl.ntree_awake': out[69],
- '_impl.nv_awake': out[70],
- '_impl.overflow': out[71],
- '_impl.qLD': out[72],
- '_impl.qLDiagInv': out[73],
- '_impl.qLU': out[74],
- 'qacc': out[75],
- 'qacc_smooth': out[76],
- 'qacc_warmstart': out[77],
- 'qfrc_actuator': out[78],
- 'qfrc_bias': out[79],
- 'qfrc_constraint': out[80],
- '_impl.qfrc_damper': out[81],
- 'qfrc_fluid': out[82],
- 'qfrc_gravcomp': out[83],
- 'qfrc_passive': out[84],
- 'qfrc_smooth': out[85],
- '_impl.qfrc_spring': out[86],
- 'qpos': out[87],
- 'qvel': out[88],
- 'sensordata': out[89],
- 'site_xmat': out[90],
- 'site_xpos': out[91],
- '_impl.solver_niter': out[92],
- '_impl.subtree_angmom': out[93],
- 'subtree_com': out[94],
- '_impl.subtree_linvel': out[95],
- '_impl.ten_J': out[96],
- 'ten_length': out[97],
- '_impl.ten_velocity': out[98],
- '_impl.ten_wrapadr': out[99],
- '_impl.ten_wrapnum': out[100],
- 'time': out[101],
- '_impl.tree_asleep': out[102],
- '_impl.tree_awake': out[103],
- '_impl.tree_island': out[104],
- '_impl.wrap_obj': out[105],
- '_impl.wrap_xpos': out[106],
- 'xanchor': out[107],
- 'xaxis': out[108],
- 'ximat': out[109],
- 'xipos': out[110],
- 'xmat': out[111],
- 'xpos': out[112],
- 'xquat': out[113],
- '_impl.contact__dim': out[114],
- '_impl.contact__dist': out[115],
- '_impl.contact__efc_address': out[116],
- '_impl.contact__elem': out[117],
- '_impl.contact__flex': out[118],
- '_impl.contact__frame': out[119],
- '_impl.contact__friction': out[120],
- '_impl.contact__geom': out[121],
- '_impl.contact__geomcollisionid': out[122],
- '_impl.contact__includemargin': out[123],
- '_impl.contact__pos': out[124],
- '_impl.contact__solimp': out[125],
- '_impl.contact__solref': out[126],
- '_impl.contact__solreffriction': out[127],
- '_impl.contact__type': out[128],
- '_impl.contact__vert': out[129],
- '_impl.contact__worldid': out[130],
- '_impl.efc__D': out[131],
- '_impl.efc__J': out[132],
- '_impl.efc__J_colind': out[133],
- '_impl.efc__J_rowadr': out[134],
- '_impl.efc__J_rownnz': out[135],
- '_impl.efc__Jqvel': out[136],
- '_impl.efc__Ma': out[137],
- '_impl.efc__aref': out[138],
- '_impl.efc__force': out[139],
- '_impl.efc__frictionloss': out[140],
- '_impl.efc__id': out[141],
- '_impl.efc__island': out[142],
- '_impl.efc__jtdaj_adr': out[143],
- '_impl.efc__jtdaj_nblock': out[144],
- '_impl.efc__jtdaj_nrow': out[145],
- '_impl.efc__margin': out[146],
- '_impl.efc__pos': out[147],
- '_impl.efc__state': out[148],
- '_impl.efc__type': out[149],
- '_impl.efc__vel': out[150],
+ '_impl.cqLD': out[20],
+ '_impl.cqacc_smooth': out[21],
+ '_impl.cqacc_warmstart': out[22],
+ '_impl.cqfrc_smooth': out[23],
+ '_impl.crb': out[24],
+ '_impl.crhs': out[25],
+ 'cvel': out[26],
+ '_impl.cx': out[27],
+ '_impl.dof_awake_ind': out[28],
+ '_impl.dof_cdof': out[29],
+ '_impl.dof_island': out[30],
+ '_impl.dof_islandid': out[31],
+ '_impl.efc_islandid': out[32],
+ '_impl.energy': out[33],
+ '_impl.face_quat': out[34],
+ '_impl.face_xpos': out[35],
+ '_impl.flex_aabb_max': out[36],
+ '_impl.flex_aabb_min': out[37],
+ '_impl.flexedge_J': out[38],
+ '_impl.flexedge_length': out[39],
+ '_impl.flexedge_velocity': out[40],
+ '_impl.flexnode_xpos': out[41],
+ '_impl.flexvert_xpos': out[42],
+ 'geom_xmat': out[43],
+ 'geom_xpos': out[44],
+ 'history': out[45],
+ '_impl.island_dofadr': out[46],
+ '_impl.island_idofadr': out[47],
+ '_impl.island_iefcadr': out[48],
+ '_impl.island_ne': out[49],
+ '_impl.island_nefc': out[50],
+ '_impl.island_nf': out[51],
+ '_impl.island_nv': out[52],
+ '_impl.light_xdir': out[53],
+ '_impl.light_xpos': out[54],
+ '_impl.map_dof2idof': out[55],
+ '_impl.map_efc2iefc': out[56],
+ '_impl.map_idof2dof': out[57],
+ '_impl.map_iefc2efc': out[58],
+ '_impl.moment_colind': out[59],
+ '_impl.moment_rowadr': out[60],
+ '_impl.moment_rownnz': out[61],
+ '_impl.nacon': out[62],
+ '_impl.nbody_awake': out[63],
+ '_impl.ncdof': out[64],
+ '_impl.ncollision': out[65],
+ '_impl.ne': out[66],
+ '_impl.nefc': out[67],
+ '_impl.nf': out[68],
+ '_impl.nidof': out[69],
+ '_impl.nisland': out[70],
+ '_impl.nl': out[71],
+ '_impl.ntree_awake': out[72],
+ '_impl.nv_awake': out[73],
+ '_impl.overflow': out[74],
+ '_impl.qLD': out[75],
+ '_impl.qLDiagInv': out[76],
+ '_impl.qLU': out[77],
+ 'qacc': out[78],
+ 'qacc_smooth': out[79],
+ 'qacc_warmstart': out[80],
+ 'qfrc_actuator': out[81],
+ 'qfrc_bias': out[82],
+ 'qfrc_constraint': out[83],
+ '_impl.qfrc_damper': out[84],
+ 'qfrc_fluid': out[85],
+ 'qfrc_gravcomp': out[86],
+ 'qfrc_passive': out[87],
+ 'qfrc_smooth': out[88],
+ '_impl.qfrc_spring': out[89],
+ 'qpos': out[90],
+ 'qvel': out[91],
+ 'sensordata': out[92],
+ 'site_xmat': out[93],
+ 'site_xpos': out[94],
+ '_impl.solver_niter': out[95],
+ '_impl.subtree_angmom': out[96],
+ 'subtree_com': out[97],
+ '_impl.subtree_linvel': out[98],
+ '_impl.ten_J': out[99],
+ 'ten_length': out[100],
+ '_impl.ten_velocity': out[101],
+ '_impl.ten_wrapadr': out[102],
+ '_impl.ten_wrapnum': out[103],
+ 'time': out[104],
+ '_impl.tree_asleep': out[105],
+ '_impl.tree_awake': out[106],
+ '_impl.tree_island': out[107],
+ '_impl.wrap_obj': out[108],
+ '_impl.wrap_xpos': out[109],
+ 'xanchor': out[110],
+ 'xaxis': out[111],
+ 'ximat': out[112],
+ 'xipos': out[113],
+ 'xmat': out[114],
+ 'xpos': out[115],
+ 'xquat': out[116],
+ '_impl.contact__dim': out[117],
+ '_impl.contact__dist': out[118],
+ '_impl.contact__efc_address': out[119],
+ '_impl.contact__elem': out[120],
+ '_impl.contact__flex': out[121],
+ '_impl.contact__frame': out[122],
+ '_impl.contact__friction': out[123],
+ '_impl.contact__geom': out[124],
+ '_impl.contact__geomcollisionid': out[125],
+ '_impl.contact__includemargin': out[126],
+ '_impl.contact__pos': out[127],
+ '_impl.contact__solimp': out[128],
+ '_impl.contact__solref': out[129],
+ '_impl.contact__solreffriction': out[130],
+ '_impl.contact__type': out[131],
+ '_impl.contact__vert': out[132],
+ '_impl.contact__worldid': out[133],
+ '_impl.efc__D': out[134],
+ '_impl.efc__J': out[135],
+ '_impl.efc__J_colind': out[136],
+ '_impl.efc__J_rowadr': out[137],
+ '_impl.efc__J_rownnz': out[138],
+ '_impl.efc__Jqvel': out[139],
+ '_impl.efc__Ma': out[140],
+ '_impl.efc__aref': out[141],
+ '_impl.efc__force': out[142],
+ '_impl.efc__frictionloss': out[143],
+ '_impl.efc__id': out[144],
+ '_impl.efc__island': out[145],
+ '_impl.efc__jtdaj_adr': out[146],
+ '_impl.efc__jtdaj_nblock': out[147],
+ '_impl.efc__jtdaj_nrow': out[148],
+ '_impl.efc__margin': out[149],
+ '_impl.efc__pos': out[150],
+ '_impl.efc__state': out[151],
+ '_impl.efc__type': out[152],
+ '_impl.efc__vel': out[153],
})
return d
diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py
index cbfcec54..194bf2d2 100644
--- a/mjx/mujoco/mjx/warp/render.py
+++ b/mjx/mujoco/mjx/warp/render.py
@@ -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,
diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py
index ce765d8b..ff0ba1e7 100644
--- a/mjx/mujoco/mjx/warp/smooth.py
+++ b/mjx/mujoco/mjx/warp/smooth.py
@@ -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,
)
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index ad11625d..0cfa5409 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -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,
diff --git a/mjx/mujoco/mjx/warp/types_test.py b/mjx/mujoco/mjx/warp/types_test.py
index 6d1bb9af..6bc1d8ae 100644
--- a/mjx/mujoco/mjx/warp/types_test.py
+++ b/mjx/mujoco/mjx/warp/types_test.py
@@ -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('')
+ 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', ""),
+ 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):