Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 960411406
Change-Id: I687b06586881d1915bab1bf8ebab6473e6ef29db

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_)
+  # Match and convert custom mujoco_warp array annotations (e.g. within tuple[...]).
-  if m:  # match custom mujoco_warp array annotation types
+  def _tuple_array(m):
-    args_str = m.group(1)
+    args = [a.strip() for a in m.group(1).split(',')]
-    args = [a.strip() for a in args_str.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)
+  ModelWarp = getattr(mjxw_types, 'ModelWarp', None)
-  OptionWarp = getattr(mjx_types, 'OptionWarp', 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):
+    if is_jax_array(mjx_types.Model, field) or is_jax_array(ModelWarp, field):
-      stage_in.append(field)
-    elif 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):
+      if is_jax_array(mjx_types.Option, sub_field) or is_jax_array(OptionWarp, sub_field):
-        stage_in.append(field)
-      elif is_jax_array(OptionWarp, sub_field):
         stage_in.append(field)
 
-  # stage_in: Data jax.Array input fields
+  # stage_in: Data/DataWarp jax.Array input fields
   for field in field_usage.data_fields:
-    if is_jax_array(mjx_types.Data, field):
+    if is_jax_array(mjx_types.Data, field) or is_jax_array(DataWarp, field):
       stage_in.append(field)
 
-  # stage_out: Data jax.Array output fields
+  # stage_out: Data/DataWarp jax.Array output fields
   for field in field_usage.data_out_fields:
-    if is_jax_array(mjx_types.Data, field):
+    if is_jax_array(mjx_types.Data, field) or is_jax_array(DataWarp, field):
       stage_out.append(field)
 
   return sorted(stage_in), sorted(stage_out)

mjx/mujoco/mjx/codegen/generate_warp_types.py

@@ -19,19 +19,17 @@ import dataclasses
 import enum
 import logging
 import typing
-from typing import Any, Callable, Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional, Set
 
 from absl import app
 from absl import flags
 from etils import epath
-import mujoco
 from mujoco.mjx.codegen import file
 import mujoco.mjx.third_party.mujoco_warp as mjwarp
 import numpy as np
 import warp as wp
-from warp._src.jax.ffi import FfiArg
 from warp import JaxCallableGraphMode as GraphMode
-
+from warp._src.jax.ffi import FfiArg
 
 _MJX_WARP_TYPES_OUT_FPATH = flags.DEFINE_string(
     'mjx_warp_types_out_path',
@@ -46,20 +44,33 @@ _MJX_TYPES_PATH = flags.DEFINE_string(
 )
 
 _DATA_SHAPE_PROPERTY_FIELD = 'cacc'
-_DUMMY_XML = """
+
-  <mujoco>
+_CLS_MAP = {
-    <worldbody>
+    'Model': mjwarp.Model,
-      <body>
+    'Data': mjwarp.Data,
-        <joint type="free"/>
+    'Option': mjwarp.Option,
-        <geom pos="0 0 0" size="0.2" type="sphere"/>
+    'Statistic': mjwarp.Statistic,
-      </body>
+}
-      <body >
+
-        <joint type="free"/>
+
-        <geom pos="0 0.3 0" size="0.11" type="sphere"/>
+def _is_array_spec(typ) -> bool:
-      </body>
+  """Check if a type annotation is a Warp array spec."""
-    </worldbody>
+  return isinstance(typ, wp.array) or type(typ).__name__ == '_ArrayAnnotation'
-  </mujoco>
+
-"""
+
+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
+  if _is_array_spec(annotation):
-      type(annotation).__name__ == '_ArrayAnnotation'):
     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 (
+  if is_tuple and _is_array_spec(typing.get_args(annotation)[0]):
-      isinstance(typing.get_args(annotation)[0], wp.array)
-      or type(typing.get_args(annotation)[0]).__name__ == '_ArrayAnnotation'
-  ):
     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."""
+  """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)
 
+  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'):
+    return meta_fields - _get_non_vmap_data_fields()
-      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}
 
-  with wp.ScopedDevice('cpu'):
+  return meta_fields
-    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}.')
 
 
 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 = {
+  cls = _CLS_MAP[cls_name]
-      'Model': mjwarp.Model,
-      'Data': mjwarp.Data,
-      'Option': mjwarp.Option,
-      'Statistic': mjwarp.Statistic,
-  }[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."""
+  """Returns Data fields that should not be vmapped."""
-  m = mujoco.MjModel.from_xml_string(_DUMMY_XML)
+  annotations = _get_annotations_recursive(dict(mjwarp.Data.__annotations__))
-  d = mujoco.MjData(m)
+  return {
-  mujoco.mj_step(m, d)
+      name
-
+      for name, type_ in annotations.items()
-  dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113)
+      if type_ in (int, float, bool) or _is_batched_field(type_) is False
-  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
 
 
 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
+  if _is_array_spec(wp_type):
-      type(wp_type).__name__ == '_ArrayAnnotation'):
     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]:
+  for cls_name, cls in _CLS_MAP.items():
-    ndim_annotations[cls.__name__] = {}
+    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]:
+  for cls_name, cls in _CLS_MAP.items():
-    batched[cls.__name__] = {}
+    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)
-
     all_annotations = _get_annotations_recursive(cls.__annotations__)
     for name, type_ in all_annotations.items():
       if not _is_ffi_compatible(type_):
         continue
-      batched[cls.__name__][name] = name in batched_fields
+      # Same batch check as mujoco_warp._src.io._mark_batched.
+      batched[cls_name][name] = _is_batched_field(type_) is True
 
   with target_fpath.open('a') as f:
     f.write('\n_BATCH_DIM = ' + _to_py_string(batched))

mjx/mujoco/mjx/codegen/trace.py

@@ -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):
+    if not isinstance(curr_node, ast.Name):
-      attr_parts.reverse()
+      return
-      full_attribute_str = '__'.join(attr_parts)
+    attr_parts.reverse()
-      in_field_info = full_attribute_str in self._mjwarp_field_info
+    full_attribute_str = '__'.join(attr_parts)
-      if in_field_info:
+    field_info = self._mjwarp_field_info.get(full_attribute_str)
-        if curr_node.id == 'm':
+    if field_info is None:
-          self.model_fields.add(full_attribute_str)
+      return
-        if curr_node.id == 'd':
+    # Classify by field_info.param_source; field names are guaranteed unique
-          self.data_fields.add(full_attribute_str)
+    # across Model and Data hierarchies by get_mjwarp_field_info.
-        if curr_node.id == 'd' and is_output:
+    if field_info.param_source == 'Model':
-          self.data_out_fields.add(full_attribute_str)
+      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 called_fn_name not in _BUILTIN_NAMES:
-      if key not in self._visited_fns:
+        key = (self._current_fpath, called_fn_name)
-        self._visited_fns.add(key)  # pyrefly: ignore[bad-argument-type]
+        if key not in self._visited_fns:
-        next_fpath = self._module_fpaths.get(
+          self._visited_fns.add(key)
-            called_fn_name, self._current_fpath
+          next_fpath = self._module_fpaths.get(
-        )
+              called_fn_name, self._current_fpath
-        self.recurse_trace(next_fpath, called_fn_name)
+          )
+          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:
+  for param in target_fn_node.args.args:
-    third_param = target_fn_node.args.args[2]
+    if param.annotation:
-    if third_param.annotation:
+      annotation_str = ast.unparse(param.annotation)
-      annotation_str = ast.unparse(third_param.annotation)
+      if 'RenderContext' in annotation_str:
-      render_context_in_caller = '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():
+  def get_fields(cls_name, base_order, sub_order):
-    if field in field_info:
+    fields = {}
-      raise AssertionError(f'Field {field} is duplicated.')
+    for field, typ in type_classes[cls_name].items():
-    if field in dataclass_map:
+      if field == 'callback':
-      for sfield, styp in type_classes[dataclass_map[field]].items():
+        continue
-        field_name = field + '__' + sfield
+      if field in dataclass_map:
-        field_info[field_name] = FieldInfo('Data', styp, (3, field_name))
+        for sfield, styp in type_classes[dataclass_map[field]].items():
-    else:
+          name = f'{field}__{sfield}'
-      field_info[field] = FieldInfo('Data', typ, (2, field))
+          fields[name] = FieldInfo(cls_name, styp, (sub_order, name))
-  return field_info
+      else:
+        fields[field] = FieldInfo(cls_name, typ, (base_order, field))
+    return fields
+
+  model_fields = get_fields('Model', 0, 1)
+  data_fields = get_fields('Data', 2, 3)
+
+  # Field names must be uniquely named across Model and Data hierarchies because
+  # _FunctionFieldUsageVisitor.add_field_usage attributes AST attribute accesses
+  # to Model or Data based solely on the field name lookup in mjwarp_field_info
+  # (via field_info.param_source) rather than inspecting base variable names.
+  #
+  # Nested sub-dataclasses (Option, Statistic in Model; Constraint, Contact in
+  # Data) are flattened with container prefixes (e.g. 'opt__timestep',
+  # 'contact__pos'). This prefixing prevents collisions between sub-dataclasses
+  # that share generic field names (such as 'pos' or 'type' in Contact and
+  # Constraint). The collision check below guarantees that no direct or
+  # flattened sub-dataclass field name is shared between Model and Data.
+  #
+  # These classified fields (model_fields, data_fields, data_out_fields) are
+  # then used in generate_warp_shim to build the FFI shim argument signatures.
+  if collisions := model_fields.keys() & data_fields.keys():
+    raise ValueError(
+        f'Field name collisions in MuJoCo Warp types: {collisions}'
+    )
+
+  return {**model_fields, **data_fields}
 
 
-def get_mjx_warp_field_info(src: str, get_cls_type_annotations) -> Dict[str, FieldInfo]:
+def get_mjx_warp_field_info(
+    src: str, get_cls_type_annotations
+) -> Dict[str, FieldInfo]:
   """Return field info for mjx/warp/types.py."""
   field_info = {}
   type_classes = get_cls_type_annotations(src)

mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py

@@ -41,6 +41,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import euler as euler
 from mujoco.mjx.third_party.mujoco_warp._src.forward import forward as forward
 from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_acceleration as fwd_acceleration
 from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_actuation as fwd_actuation
+from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_kinematics as fwd_kinematics
 from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_position as fwd_position
 from mujoco.mjx.third_party.mujoco_warp._src.forward import fwd_velocity as fwd_velocity
 from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit

mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py

@@ -18,6 +18,29 @@ from functools import lru_cache
 import warp as wp
 
 
+@wp.func
+def solve_search_sums(grad: float, solution: float):
+  return wp.vec2(solution * solution, grad * solution)
+
+
+@lru_cache(maxsize=None)
+def _create_newton_decrement_func(matrix_size_static: int, vector_size_static: int):
+  @wp.func
+  def newton_decrement_func(
+    # In:
+    solution_tile: wp.tile[float, matrix_size_static, 1],
+    b: wp.array2d[float],
+    # Out:
+    search_out: wp.array2d[float],
+  ):
+    grad_tile = wp.tile_load(b, shape=(vector_size_static, 1), offset=(0, 0), bounds_check=False)
+    active_solution = wp.tile_view(solution_tile, shape=(vector_size_static, 1), offset=(0, 0))
+    wp.tile_store(search_out, wp.tile_map(wp.mul, active_solution, -1.0), bounds_check=False)
+    return wp.tile_reduce(wp.add, wp.tile_map(solve_search_sums, grad_tile, active_solution))[0]
+
+  return newton_decrement_func
+
+
 @lru_cache(maxsize=None)
 def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
   @wp.func
@@ -93,7 +116,109 @@ def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_st
 
 
 @lru_cache(maxsize=None)
-def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int):
+def _create_blocked_cholesky_augmented_factorize_solve_func(
+  block_size: int,
+  matrix_size_static: int,
+  with_newton_decrement: bool,
+  vector_size_static: int,
+):
+  WITH_NEWTON_DECREMENT = with_newton_decrement
+  border_size = block_size - 1
+
+  @wp.func
+  def blocked_cholesky_augmented_factorize_solve_func(
+    # In:
+    A: wp.array2d[float],
+    b: wp.array2d[float],
+    matrix_size: int,
+    # Out:
+    U_out: wp.array2d[float],
+    result_out: wp.array2d[float],
+  ):
+    """Factor A with b as an augmented border and reuse it as the forward solution."""
+    rhs_tile = wp.tile_zeros(shape=(matrix_size_static, 1), dtype=float, storage="shared")
+
+    for k in range(0, matrix_size, block_size):
+      end = k + block_size
+      input_rhs = wp.tile_load(b, shape=(block_size, 1), offset=(k, 0), storage="shared", bounds_check=False)
+      A_kk_tile = wp.tile_load(
+        A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
+      )
+      input_diagonal_rhs = wp.tile_view(input_rhs, shape=(border_size, 1), offset=(0, 0))
+      if end == matrix_size:
+        wp.tile_assign(A_kk_tile, input_diagonal_rhs, offset=(0, border_size))
+        # Keep the augmented rank-one correction below float32 precision.
+        A_kk_tile[border_size, border_size] = 1.0e30
+
+      for j in range(0, k, block_size):
+        U_block = wp.tile_load(
+          U_out, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+        )
+        wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
+
+      wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
+      diagonal_border = wp.tile_view(A_kk_tile, shape=(border_size, 1), offset=(0, border_size))
+      if end == matrix_size:
+        wp.tile_assign(rhs_tile, diagonal_border, offset=(k, 0))
+      wp.tile_store(U_out, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
+
+      for i in range(end, matrix_size, block_size):
+        A_ki_tile = wp.tile_load(
+          A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
+        )
+        if i + block_size == matrix_size:
+          wp.tile_assign(A_ki_tile, input_rhs, offset=(0, border_size))
+
+        for j in range(0, k, block_size):
+          U_jk_tile = wp.tile_load(
+            U_out, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+          )
+          U_ji_tile = wp.tile_load(
+            U_out, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+          )
+          wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
+
+        wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
+        panel_border = wp.tile_view(A_ki_tile, shape=(block_size, 1), offset=(0, border_size))
+        if i + block_size == matrix_size:
+          wp.tile_assign(rhs_tile, panel_border, offset=(k, 0))
+        wp.tile_store(U_out, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
+
+    for i in range(matrix_size - block_size, -1, -block_size):
+      i_end = i + block_size
+      tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
+      for j in range(i_end, matrix_size, block_size):
+        U_tile = wp.tile_load(
+          U_out, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
+        x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
+
+      U_tile = wp.tile_load(
+        U_out, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
+
+    sums = wp.vec2(0.0)
+    if wp.static(WITH_NEWTON_DECREMENT):
+      sums = wp.static(_create_newton_decrement_func(matrix_size_static, vector_size_static))(rhs_tile, b, result_out)
+    else:
+      wp.tile_store(result_out, rhs_tile, offset=(0, 0), bounds_check=False)
+
+    return sums
+
+  return blocked_cholesky_augmented_factorize_solve_func
+
+
+@lru_cache(maxsize=None)
+def _create_blocked_cholesky_solve_func(
+  block_size: int,
+  matrix_size_static: int,
+  with_newton_decrement: bool,
+  vector_size_static: int,
+):
+  WITH_NEWTON_DECREMENT = with_newton_decrement
+
   @wp.func
   def blocked_cholesky_solve_func(
     # In:
@@ -101,7 +226,7 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
     b: wp.array2d[float],
     matrix_size: int,
     # Out:
-    x: wp.array2d[float],
+    result_out: wp.array2d[float],
   ):
     """Block Cholesky solve.
 
@@ -141,6 +266,30 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
 
       wp.tile_upper_solve_inplace(U_tile, tmp_tile)
 
-    wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
+    sums = wp.vec2(0.0)
+    if wp.static(WITH_NEWTON_DECREMENT):
+      sums = wp.static(_create_newton_decrement_func(matrix_size_static, vector_size_static))(rhs_tile, b, result_out)
+    else:
+      wp.tile_store(result_out, rhs_tile, offset=(0, 0), bounds_check=False)
+
+    return sums
 
   return blocked_cholesky_solve_func
+
+
+def create_blocked_cholesky_augmented_factorize_solve_func(block_size: int, matrix_size_static: int):
+  return _create_blocked_cholesky_augmented_factorize_solve_func(block_size, matrix_size_static, False, 0)
+
+
+def create_blocked_cholesky_augmented_factorize_solve_newton_func(
+  block_size: int, matrix_size_static: int, vector_size_static: int
+):
+  return _create_blocked_cholesky_augmented_factorize_solve_func(block_size, matrix_size_static, True, vector_size_static)
+
+
+def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int):
+  return _create_blocked_cholesky_solve_func(block_size, matrix_size_static, False, 0)
+
+
+def create_blocked_cholesky_solve_newton_func(block_size: int, matrix_size_static: int, vector_size_static: int):
+  return _create_blocked_cholesky_solve_func(block_size, matrix_size_static, True, vector_size_static)

mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py

@@ -28,7 +28,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @event_scope

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py

@@ -52,7 +52,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
 # TODO(team): improve compile time to enable backward pass
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float)
 mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
@@ -172,7 +172,7 @@ def ccd_hfield_kernel_builder(
   """Kernel builder for heightfield CCD collisions (no multiccd args)."""
 
   # runs convex collision on a set of geom pairs to recover contact info
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def ccd_hfield_kernel(
     # Model:
     opt_ccd_tolerance: wp.array[float],
@@ -474,6 +474,9 @@ def ccd_hfield_kernel_builder(
             epa_pr,
             epa_norm2,
             epa_horizon,
+            wp.static(warn_overflow),
+            worldid,
+            overflow_out,
           )
 
           if ncontact == 0:
@@ -833,6 +836,9 @@ def ccd_kernel_builder(
         epa_pr_in[ccdid],
         epa_norm2_in[ccdid],
         epa_horizon_in[ccdid],
+        wp.static(warn_overflow),
+        worldid,
+        overflow_out,
       )
 
     if dist >= gap and not is_collision_sensor:
@@ -946,7 +952,7 @@ def ccd_kernel_builder(
       )
 
   # runs convex collision on a set of geom pairs to recover contact info (non-heightfield)
-  @wp.kernel(module="unique", enable_backward=False, launch_bounds=(block_dim, _CCD_MIN_BLOCKS))
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False, launch_bounds=(block_dim, _CCD_MIN_BLOCKS))
   def ccd_kernel(
     # Model:
     opt_ccd_tolerance: wp.array[float],

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py

@@ -28,7 +28,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.struct

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py

@@ -41,7 +41,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 # Corresponding table to MuJoCo's mjCOLLISIONFUNC table in engine_collision_driver.c
 MJ_COLLISION_TABLE = {
@@ -373,7 +373,7 @@ def _add_geom_pair(
 
 @cache_kernel
 def _sap_project(opt_broadphase: int):
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def sap_project(
     # Model:
     ngeom: int,
@@ -430,7 +430,7 @@ def _sap_broadphase(
   enable_sleep: bool = False,
   incremental: bool = False,
 ):
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def kernel(
     # Model:
     ngeom: int,
@@ -538,7 +538,7 @@ def _sap_broadphase(
 
 @cache_kernel
 def _segmented_sort(tile_size: int):
-  @wp.kernel(module="unique")
+  @wp.kernel(module="unique", grid_stride=False)
   def segmented_sort(
     # In:
     projection_lower_in: wp.array2d[float],
@@ -691,7 +691,7 @@ def _nxn_broadphase(
   enable_sleep: bool = False,
   incremental: bool = False,
 ):
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def kernel(
     # Model:
     geom_type: wp.array[int],

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py

@@ -36,7 +36,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import OverflowType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.func
@@ -796,6 +796,7 @@ def _collide_mesh_triangle(
   epa_horizon: wp.array[int],
   tolerance: float,
   ccd_iterations: int,
+  warn_overflow: bool,
   # Data out:
   overflow_out: wp.array[int],
   # Out:
@@ -861,6 +862,9 @@ def _collide_mesh_triangle(
       epa_pr,
       epa_norm2,
       epa_horizon,
+      warn_overflow,
+      worldid,
+      overflow_out,
     )
 
     if ncontact > 0 and dist < margin + tri_radius:
@@ -1285,7 +1289,7 @@ def _plane_vertex(
   return True, dist, contact_pos, nrm_out
 
 
-@wp.kernel(module="unique", enable_backward=False)
+@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
 def _flex_internal_collisions_detect(
   # Model:
   nflex: int,
@@ -1394,7 +1398,7 @@ def _flex_internal_collisions_detect(
     )
 
 
-@wp.kernel(module="unique", enable_backward=False)
+@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
 def _flex_tet_internal_collisions_detect(
   # Model:
   nflex: int,
@@ -1840,6 +1844,7 @@ def _flex_selfcollision_narrowphase(
   # Model:
   nflex: int,
   opt_ccd_tolerance: wp.array[float],
+  opt_warn_overflow: bool,
   flex_dim: wp.array[int],
   flex_vertadr: wp.array[int],
   flex_elemadr: wp.array[int],
@@ -2015,6 +2020,9 @@ def _flex_selfcollision_narrowphase(
       epa_pr_out[pairid],
       epa_norm2_out[pairid],
       epa_horizon_out[pairid],
+      opt_warn_overflow,
+      worldid,
+      overflow_out,
     )
 
     phys_dist = dist
@@ -2055,11 +2063,12 @@ def _flex_selfcollision_narrowphase(
       )
 
 
-@wp.kernel(module="unique", enable_backward=False)
+@wp.kernel(module="unique", enable_backward=False, grid_stride=False)
 def _flex_active_element_collisions_detect(
   # Model:
   nflex: int,
   opt_ccd_tolerance: wp.array[float],
+  opt_warn_overflow: bool,
   flex_selfcollide: wp.array[int],
   flex_dim: wp.array[int],
   flex_vertadr: wp.array[int],
@@ -2234,6 +2243,9 @@ def _flex_active_element_collisions_detect(
         epa_pr_out[unique_thread_id],
         epa_norm2_out[unique_thread_id],
         epa_horizon_out[unique_thread_id],
+        opt_warn_overflow,
+        worldid,
+        overflow_out,
       )
 
       phys_dist = dist
@@ -2445,6 +2457,7 @@ def _flex_narrowphase_unified(
         epa_horizon[ccdid],
         tolerance,
         ccd_iterations,
+        opt_warn_overflow,
         overflow_out,
         cand_dist_out,
         cand_pos_out,
@@ -2507,6 +2520,9 @@ def _flex_narrowphase_unified(
           epa_pr[ccdid],
           epa_norm2[ccdid],
           epa_horizon[ccdid],
+          opt_warn_overflow,
+          worldid,
+          overflow_out,
         )
 
         if ncontact > 0 and dist < margin + tri_radius:
@@ -3650,6 +3666,7 @@ def flex_collision(m: Model, d: Data, ctx):
         inputs=[
           m.nflex,
           m.opt.ccd_tolerance,
+          m.opt.warn_overflow,
           m.flex_dim,
           m.flex_vertadr,
           m.flex_elemadr,
@@ -3716,6 +3733,7 @@ def flex_collision(m: Model, d: Data, ctx):
         inputs=[
           m.nflex,
           m.opt.ccd_tolerance,
+          m.opt.warn_overflow,
           m.flex_selfcollide,
           m.flex_dim,
           m.flex_vertadr,

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py

@@ -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
+  lmbda = wp.vec4(1.0, 0.0, 0.0, 0.0)  # barycentric coordinates
-  tol2 = tolerance * tolerance
+
-  epsilon = wp.where(is_discrete, 0.0, 0.5 * tol2)
+  # 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)
+    if xnorm < min_norm or wp.abs(xnorm_prev - xnorm) < min_tol:
-    # TODO(kbayes): determine new constant here
-    if xnorm2 < tol2 or wp.abs(xnorm2_old - xnorm2) < tol2:
       break
-    xnorm2_old = xnorm2
-    dir_neg = x_k / wp.sqrt(xnorm2)
 
-    # compute kth support point in geom1
+    # compute the support point with direction tuning
-    sp = support(geom1, geomtype1, -dir_neg)
+    sp1, sp2 = _gjk_support(geom1, geom2, geomtype1, geomtype2, x_k, xnorm, simplex, n, is_discrete)
-    simplex1[n] = sp.point
+    simplex1[n] = sp1.point
-    geom1.index = sp.cached_index
+    geom1.index = sp1.cached_index
-    simplex_index1[n] = sp.vertex_index
+    simplex_index1[n] = sp1.vertex_index
 
-    # compute kth support point in geom2
+    simplex2[n] = sp2.point
-    sp = support(geom2, geomtype2, dir_neg)
+    geom2.index = sp2.cached_index
-    simplex2[n] = sp.point
+    simplex_index2[n] = sp2.vertex_index
-    geom2.index = sp.cached_index
-    simplex_index2[n] = sp.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
+    if dist[idx] > MINVAL:
-    return pt, GJKResult()
+      pt.status = 12
+      return pt, GJKResult()
+
+    # fallback to closest face
+    pt.status = -1
+    if idx == 0:
+      return pt, _replace_simplex3(pt, 0, 1, 2)
+    elif idx == 1:
+      return pt, _replace_simplex3(pt, 0, 3, 1)
+    elif idx == 2:
+      return pt, _replace_simplex3(pt, 0, 2, 3)
+    else:
+      return pt, _replace_simplex3(pt, 3, 2, 1)
 
   # set polytope counts
   pt.nvert = 4
@@ -1251,6 +1315,10 @@ def _epa(
   geomtype1: int,
   geomtype2: int,
   is_discrete: bool,
+  warn_overflow: bool,
+  worldid: int,
+  # Data out:
+  overflow_out: wp.array[int],
 ) -> Tuple[float, wp.vec3, wp.vec3, int]:
   """Recover penetration data from two geoms in contact given an initial polytope."""
   upper = FLOAT_MAX
@@ -1320,7 +1388,9 @@ def _epa(
     pt.nhorizon = _add_edge(pt, face[1], face[2])
     pt.nhorizon = _add_edge(pt, face[2], face[0])
     if pt.nhorizon == -1:
-      wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+      if warn_overflow:
+        wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+      wp.atomic_or(overflow_out, worldid, OverflowType.EPA_HORIZON)
       idx = -1
       break
 
@@ -1337,7 +1407,9 @@ def _epa(
         pt.nhorizon = _add_edge(pt, face[1], face[2])
         pt.nhorizon = _add_edge(pt, face[2], face[0])
         if pt.nhorizon == -1:
-          wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+          if warn_overflow:
+            wp.printf("Warning: EPA horizon = %d isn't large enough.\n", pt.horizon.shape[0])
+          wp.atomic_or(overflow_out, worldid, OverflowType.EPA_HORIZON)
           idx = -1
           break
 
@@ -2302,6 +2374,8 @@ def gjk_phase(
   if size1 + size2 > 0.0:
     cutoff += full_margin1 + full_margin2
     result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
+    geom1.index = result.index1
+    geom2.index = result.index2
 
     # shallow penetration, inflate contact
     if result.dist > tolerance:
@@ -2318,6 +2392,8 @@ def gjk_phase(
     cutoff -= full_margin1 + full_margin2
 
   result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
+  geom1.index = result.index1
+  geom2.index = result.index2
 
   # no penetration depth to recover
   if result.dist > tolerance or result.dim < 2:
@@ -2342,6 +2418,10 @@ def epa_phase(
   face_pr: wp.array[wp.vec3],
   face_norm2: wp.array[float],
   horizon: wp.array[int],
+  warn_overflow: bool,
+  worldid: int,
+  # Data out:
+  overflow_out: wp.array[int],
 ) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
   """Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
   pt = Polytope()
@@ -2413,7 +2493,9 @@ def epa_phase(
     return result.dist, 1, result.x1, result.x2, -1
 
   is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
-  dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
+  dist, x1, x2, idx = _epa(
+    tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete, warn_overflow, worldid, overflow_out
+  )
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
 
@@ -2447,6 +2529,10 @@ def ccd(
   face_pr: wp.array[wp.vec3],
   face_norm2: wp.array[float],
   horizon: wp.array[int],
+  warn_overflow: bool,
+  worldid: int,
+  # Data out:
+  overflow_out: wp.array[int],
 ) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
   """General convex collision detection via GJK/EPA."""
   needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
@@ -2468,4 +2554,7 @@ def ccd(
     face_pr,
     face_norm2,
     horizon,
+    warn_overflow,
+    worldid,
+    overflow_out,
   )

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive_core.py

@@ -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
+from mujoco.mjx.third_party.mujoco_warp._src.math import closest_segment_point
-MJ_MAXVAL = 1e10
+from mujoco.mjx.third_party.mujoco_warp._src.math import closest_segment_to_segment_points
-
+from mujoco.mjx.third_party.mujoco_warp._src.math import normalize_with_norm
+from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 
 wp.set_module_options({"enable_backward": False})
 
 
-@wp.func
-def safe_div(x: Any, y: Any) -> Any:
-  return x / wp.where(y != 0.0, y, MJ_MINVAL)
-
-
-@wp.func
-def normalize_with_norm(x: Any):
-  norm = wp.length(x)
-  if norm == 0.0:
-    return x, 0.0
-  return x / norm, norm
-
-
-@wp.func
-def closest_segment_point(a: wp.vec3, b: wp.vec3, pt: wp.vec3) -> wp.vec3:
-  """Returns the closest point on the a-b line segment to a point pt."""
-  ab = b - a
-  t = wp.dot(pt - a, ab) / (wp.dot(ab, ab) + 1e-6)
-  return a + wp.clamp(t, 0.0, 1.0) * ab
-
-
-@wp.func
-def closest_segment_point_and_dist(a: wp.vec3, b: wp.vec3, pt: wp.vec3) -> Tuple[wp.vec3, float]:
-  """Returns closest point on the line segment and the distance squared."""
-  closest = closest_segment_point(a, b, pt)
-  dist = wp.dot((pt - closest), (pt - closest))
-  return closest, dist
-
-
-@wp.func
-def closest_segment_to_segment_points(a0: wp.vec3, a1: wp.vec3, b0: wp.vec3, b1: wp.vec3) -> Tuple[wp.vec3, wp.vec3]:
-  """Returns closest points between two line segments."""
-  dir_a, len_a = normalize_with_norm(a1 - a0)
-  dir_b, len_b = normalize_with_norm(b1 - b0)
-
-  half_len_a = len_a * 0.5
-  half_len_b = len_b * 0.5
-  a_mid = a0 + dir_a * half_len_a
-  b_mid = b0 + dir_b * half_len_b
-
-  trans = a_mid - b_mid
-
-  dira_dot_dirb = wp.dot(dir_a, dir_b)
-  dira_dot_trans = wp.dot(dir_a, trans)
-  dirb_dot_trans = wp.dot(dir_b, trans)
-  denom = 1.0 - dira_dot_dirb * dira_dot_dirb
-
-  orig_t_a = (-dira_dot_trans + dira_dot_dirb * dirb_dot_trans) / (denom + 1e-6)
-  orig_t_b = dirb_dot_trans + orig_t_a * dira_dot_dirb
-  t_a = wp.clamp(orig_t_a, -half_len_a, half_len_a)
-  t_b = wp.clamp(orig_t_b, -half_len_b, half_len_b)
-
-  best_a = a_mid + dir_a * t_a
-  best_b = b_mid + dir_b * t_b
-
-  new_a, d1 = closest_segment_point_and_dist(a0, a1, best_b)
-  new_b, d2 = closest_segment_point_and_dist(b0, b1, best_a)
-  if d1 < d2:
-    return new_a, best_b
-  return best_a, new_b
-
-
 class vec8f(wp.types.vector(length=8, dtype=wp.float32)):
   pass
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py

@@ -34,7 +34,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec_pluginattr
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.struct

mjx/mujoco/mjx/third_party/mujoco_warp/_src/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_actlimited,
-          m.actuator_forcerange,
           m.actuator_actrange,
+          m.actuator_actearly,
+          m.actuator_forcelimited,
+          m.actuator_forcerange,
           d.act,
           d.ctrl,
           d.act_dot,

mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py

@@ -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):
+def fwd_kinematics(m: Model, d: Data):
-  """Position-dependent computations.
+  """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_dynprm: wp.array2d[vec10],
-  actuator_forcelimited: wp.array[bool],
+  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_actlimited,
-      m.actuator_ctrlrange,
-      m.actuator_forcerange,
       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:
+  _energy_pos(m, d)
-    if m.sensor_e_potential == 0:  # not computed by sensor
-      sensor.energy_pos(m, d)
-  else:
-    d.energy.zero_()
 
   fwd_velocity(m, d)
   sensor.sensor_vel(m, d)
-
+  _energy_vel(m, d)
-  if energy:
-    if m.sensor_e_kinetic == 0:  # not computed by sensor
-      sensor.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:
+  _energy_pos(m, d)
-    if m.sensor_e_potential == 0:  # not computed by sensor
-      sensor.energy_pos(m, d)
-  else:
-    d.energy.zero_()
 
   fwd_velocity(m, d)
   sensor.sensor_vel(m, d)
 
-  if energy:
+  _energy_vel(m, d)
-    if m.sensor_e_kinetic == 0:  # not computed by sensor
-      sensor.energy_vel(m, d)
 
   if not (m.opt.disableflags & DisableBit.ACTUATION):
     if m.callback.control:

mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py

@@ -20,7 +20,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.func

mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py

@@ -29,7 +29,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.kernel
@@ -119,6 +119,13 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
   smooth.solve_m(m, d, qacc, qfrc)
 
 
+@wp.kernel
+def _zero_qfrc_constraint_nefc(nefc_in: wp.array[int], qfrc_constraint_out: wp.array2d[float]):
+  worldid, dofid = wp.tid()
+  if nefc_in[worldid] == 0:
+    qfrc_constraint_out[worldid, dofid] = 0.0
+
+
 def inv_constraint(m: Model, d: Data):
   """Inverse constraint solver."""
   # no constraints
@@ -126,6 +133,14 @@ def inv_constraint(m: Model, d: Data):
     d.qfrc_constraint.zero_()
     return
 
+  if m.is_sparse:
+    wp.launch(
+      _zero_qfrc_constraint_nefc,
+      dim=(d.nworld, m.nv),
+      inputs=[d.nefc],
+      outputs=[d.qfrc_constraint],
+    )
+
   ctx = solver.create_inverse_context(m, d)
   solver.init_context(m, d, ctx, grad=False)
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -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
+  Blocks use scalar Cholesky through six DOFs, tile Cholesky through M_BLOCK_DENSE_MAX, and sparse
-  categories by coupling and size:
+  LDL beyond that. Compact diagonal blocks also use the scalar path without allocating a factor.
-    - 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
   """
   nv = mjm.nv
   blocks = _m_blocks(mjm)
-  allow_dense = _m_allow_dense(mjm)
+  dof_adr = np.full(nv, types.Q_LD_BLOCK_SPARSE, dtype=np.int32)
-  rownnz = mjm.M_rownnz
+  scalar_tiles = {}
-  dof_adr = np.zeros(nv, dtype=np.int32)
+  gather_tiles = {}
-  dof_dense = np.zeros(nv, dtype=np.int32)
-  dof_simple = np.zeros(nv, dtype=np.int32)
-  dense_blocks = []
   off = 0
-  has_sparse = False
   for start, size in blocks:
-    coupled = bool(np.max(rownnz[start : start + size]) > 1)
+    last = start + size - 1
-    if not coupled:
+    madr = int(mjm.M_rowadr[start])
-      dof_simple[start : start + size] = 1
+    nnz = int(mjm.M_rowadr[last] + mjm.M_rownnz[last] - madr)
-    elif allow_dense and size <= types.M_BLOCK_DENSE_MAX:
+    compact = nnz == size
-      dense_blocks.append((start, 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:
+  for starts in scalar_tiles.values():
-      has_sparse = True
+    starts.sort(key=lambda start: dof_adr[start] >= 0)
   return {
     "total": off,
     "dof_adr": dof_adr,
-    "blocks": blocks,
+    "scalar_tiles": scalar_tiles,
-    "dense_blocks": dense_blocks,
+    "gather_tiles": gather_tiles,
-    "dof_dense": dof_dense,
+    "has_sparse": bool(np.any(dof_adr == types.Q_LD_BLOCK_SPARSE)),
-    "dof_simple": dof_simple,
-    "has_dense": len(dense_blocks) > 0,
-    "has_simple": bool(dof_simple.any()),
-    "has_sparse": has_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
+      ny_g = cy * order_abs + 1
-      nz_g = cz * order + 1
+      nz_g = cz * order_abs + 1
 
-      node_bodies = [
+      if order < 0:
-        mjm.flex_nodebodyid[nstart + (ci * order + li) * ny_g * nz_g + (cj * order + lj) * nz_g + (ck * order + lk)]
+        # Shell mode: 2D bilinear quad
-        for li in range(order + 1)
+        npc = (order_abs + 1) * (order_abs + 1)
-        for lj in range(order + 1)
+        node_bodies = []
-        for lk in range(order + 1)
+        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
+  # Per-block scalar/tile/sparse layout (see m_block_layout).
-  # size; a model may use both paths at once (e.g. one large tree + many small free joints).
   _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 = {}
+  scalar_tiles = [
-  for start, size in _lay["dense_blocks"]:
+    types.TileSet(
-    tiles.setdefault(size, []).append(start)
+      adr=wp.array(_lay["scalar_tiles"][size], dtype=int),
-  m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
+      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
+  # Group sparse LDL updates by tree depth. Block-path DOFs never touch the LDL region.
-  # dofs are included; dense blocks use the packed Cholesky and never touch the LDL region.
+  sparse_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
-  qLD_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]:
+    if _lay["dof_adr"][k] != types.Q_LD_BLOCK_SPARSE:
-      continue  # dense block: handled by the packed Cholesky, not the LDL factor
+      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):
+  for level in sorted(sparse_updates):
-    all_updates_flat.extend(qLD_updates[level])
+    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):
+def _get_padded_sizes(nv: int, njmax: int, is_sparse: bool, tile_size: int, augment_cholesky: bool = False):
-  # 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 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_set_nnz(mjm, cell_bodies | {gb}))
-          max_contact_nnz = max(max_contact_nnz, _body_pair_nnz(mjm, fb, 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, sizes, mjd)
-
-  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, 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 block Cholesky | nC LDL region]. Block factors store their upper Cholesky
-  # qLD = [packed dense-block Cholesky | nC LDL region]. Dense blocks 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"]:
+    for size, starts in list(lay["scalar_tiles"].items()) + list(lay["gather_tiles"].items()):
-      off = lay["dof_adr"][start]
+      for start in starts:
-      blk = Mfull[start : start + size, start : start + size]
+        off = lay["dof_adr"][start]
-      if blk.any():
+        if off < 0:
-        qLD[off : off + size * size] = np.linalg.cholesky(blk).T.reshape(-1)
+          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"]:
+  if _lay["scalar_tiles"] or _lay["gather_tiles"]:
-    # d.qLD is not MuJoCo's LDL: dense blocks are a packed Cholesky and simple blocks are factored
+    # Block factors do not use MuJoCo's LDL representation.
-    # into qLDiagInv (their LDL slots are never written). Recompute the LDL factor from M.
     mujoco.mj_factorM(mjm, result)
   else:
     # Pure sparse: qLD is exactly MuJoCo's nC LDL factor.
@@ -2227,7 +2439,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
   """
   sleep_enabled = bool(m.opt.enableflags & types.EnableBit.SLEEP)
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def reset_xfrc_applied(reset_in: wp.array[bool], xfrc_applied_out: wp.array2d[wp.spatial_vector]):
     worldid, bodyid, elemid = wp.tid()
 
@@ -2237,7 +2449,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
 
     xfrc_applied_out[worldid, bodyid][elemid] = 0.0
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def reset_M(reset_in: wp.array[bool], M_out: wp.array2d[float]):
     worldid, elemid = wp.tid()
 
@@ -2247,7 +2459,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
 
     M_out[worldid, elemid] = 0.0
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def reset_nworld(
     # Model:
     nq: int,
@@ -2329,7 +2541,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
       userdata_out[worldid, i] = 0.0
     overflow_out[worldid] = 0
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def reset_mocap(
     # Model:
     body_mocapid: wp.array[int],
@@ -2353,7 +2565,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
       mocap_pos_out[worldid, mocapid] = body_pos[worldid % body_pos.shape[0], bodyid]
       mocap_quat_out[worldid, mocapid] = body_quat[worldid % body_quat.shape[0], bodyid]
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def reset_contact(
     # Data in:
     nacon_in: wp.array[int],
@@ -2412,7 +2624,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     contact_type_out[conid] = 0
     contact_geomcollisionid_out[conid] = 0
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def reset_sleep(
     # Model:
     nv: int,
@@ -3051,14 +3263,14 @@ def _compute_dof_M0(
 @wp.kernel
 def _resolve_dampratio(
   actuator_biastype: wp.array[int],
-  actuator_gainprm: wp.array2d[types.vec10f],
+  actuator_gainprm: wp.array2d[types.vec10],
   moment_rownnz_in: wp.array2d[int],
   moment_rowadr_in: wp.array2d[int],
   moment_colind_in: wp.array2d[int],
   actuator_moment_in: wp.array2d[float],
   dof_M0_in: wp.array2d[float],
   nv: int,
-  actuator_biasprm: wp.array2d[types.vec10f],
+  actuator_biasprm: wp.array2d[types.vec10],
 ):
   worldid, actid = wp.tid()
   biastype = actuator_biastype[actid]
@@ -3149,12 +3361,13 @@ def set_const_fixed(m: types.Model, d: types.Data):
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
   """
-  wp.launch(_init_subtreemass, dim=(d.nworld, m.nbody), inputs=[m.body_mass], outputs=[m.body_subtreemass])
+  nworld_subtreemass = m.body_subtreemass.shape[0]
+  wp.launch(_init_subtreemass, dim=(nworld_subtreemass, m.nbody), inputs=[m.body_mass], outputs=[m.body_subtreemass])
   for i in reversed(range(len(m.body_tree))):
     body_tree = m.body_tree[i]
     wp.launch(
       _accumulate_subtreemass,
-      dim=(d.nworld, body_tree.size),
+      dim=(nworld_subtreemass, body_tree.size),
       inputs=[m.body_parentid, m.body_subtreemass, body_tree],
     )
 
@@ -3198,16 +3411,16 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
   # Compute meaninertia from M diagonal at qpos0
   wp.launch(
     _compute_meaninertia,
-    dim=d.nworld,
+    dim=m.stat.meaninertia.shape[0],
     inputs=[m.nv, m.M_rownnz, m.M_rowadr, d.M],
     outputs=[m.stat.meaninertia],
   )
 
-  wp.launch(_copy_tendon_length0, dim=(d.nworld, m.ntendon), inputs=[d.ten_length], outputs=[m.tendon_length0])
+  wp.launch(_copy_tendon_length0, dim=(m.tendon_length0.shape[0], m.ntendon), inputs=[d.ten_length], outputs=[m.tendon_length0])
 
   wp.launch(
     _compute_eq_data0,
-    dim=(d.nworld, m.neq),
+    dim=(m.eq_data.shape[0], m.neq),
     inputs=[m.eq_type, m.eq_obj1id, m.eq_obj2id, m.eq_objtype, d.xpos, d.xquat, d.xmat],
     outputs=[m.eq_data],
   )
@@ -3229,7 +3442,7 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
 
     wp.launch(
       _finalize_dof_invweight0,
-      dim=(d.nworld, m.nv),
+      dim=(m.dof_invweight0.shape[0], m.nv),
       inputs=[m.dof_jntid, m.jnt_type, m.jnt_dofadr, dof_A_diag],
       outputs=[m.dof_invweight0],
     )
@@ -3272,7 +3485,7 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
 
     wp.launch(
       _finalize_body_invweight0,
-      dim=(d.nworld, m.nbody),
+      dim=(m.body_invweight0.shape[0], m.nbody),
       inputs=[m.body_weldid, body_A_diag],
       outputs=[m.body_invweight0],
     )
@@ -3295,21 +3508,23 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
       smooth.solve_m(m, d, ten_result_vec, ten_J_vec)
       wp.launch(
         _compute_tendon_dot_product,
-        dim=d.nworld,
+        dim=m.tendon_invweight0.shape[0],
         inputs=[m.ten_J_rownnz, m.ten_J_rowadr, m.ten_J_colind, tenid, d.ten_J, ten_result_vec],
         outputs=[m.tendon_invweight0],
       )
 
+  nworld_cam = np.max([m.cam_pos0.shape[0], m.cam_poscom0.shape[0], m.cam_mat0.shape[0]])
   wp.launch(
     _compute_cam_pos0,
-    dim=(d.nworld, m.ncam),
+    dim=(nworld_cam, m.ncam),
     inputs=[m.cam_bodyid, m.cam_targetbodyid, d.cam_xpos, d.cam_xmat, d.xpos, d.subtree_com],
     outputs=[m.cam_pos0, m.cam_poscom0, m.cam_mat0],
   )
 
+  nworld_light = np.max([m.light_pos0.shape[0], m.light_poscom0.shape[0], m.light_dir0.shape[0]])
   wp.launch(
     _compute_light_pos0,
-    dim=(d.nworld, m.nlight),
+    dim=(nworld_light, m.nlight),
     inputs=[m.light_bodyid, m.light_targetbodyid, d.light_xpos, d.light_xdir, d.xpos, d.subtree_com],
     outputs=[m.light_pos0, m.light_poscom0, m.light_dir0],
   )
@@ -3327,7 +3542,9 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
         outputs=[act_moment_vec],
       )
       smooth.solve_m(m, d, act_result_vec, act_moment_vec)
-      wp.launch(_compute_actuator_acc0, dim=d.nworld, inputs=[actid, m.nv, act_result_vec], outputs=[m.actuator_acc0])
+      wp.launch(
+        _compute_actuator_acc0, dim=m.actuator_acc0.shape[0], inputs=[actid, m.nv, act_result_vec], outputs=[m.actuator_acc0]
+      )
 
   # resolve dampratio: compute dof_M0, then convert dampratio to damping
   if m.nu > 0 and m.nv > 0:
@@ -3340,7 +3557,7 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
     )
     wp.launch(
       _resolve_dampratio,
-      dim=(d.nworld, m.nu),
+      dim=(m.actuator_biasprm.shape[0], m.nu),
       inputs=[
         m.actuator_biastype,
         m.actuator_gainprm,
@@ -3388,7 +3605,7 @@ def set_const_spring(m: types.Model, d: types.Data, restore: bool = True):
 
   wp.launch(
     _resolve_tendon_lengthspring,
-    dim=(d.nworld, m.ntendon),
+    dim=(m.tendon_lengthspring.shape[0], m.ntendon),
     inputs=[d.ten_length],
     outputs=[m.tendon_lengthspring],
   )
@@ -3701,6 +3918,7 @@ def create_render_context(
   render_depth: list[bool] | bool | None = None,
   render_seg: list[bool] | bool | None = None,
   use_textures: bool = True,
+  use_fast_math: bool = True,
   use_shadows: bool = False,
   use_ambient_lighting: bool = True,
   enabled_geom_groups: list[int] = [0, 1, 2],
@@ -3726,6 +3944,7 @@ def create_render_context(
     render_seg: Whether to render segmentation (per-pixel object ID/type pairs).
       If None, uses the MuJoCo model values.
     use_textures: Whether to use textures.
+    use_fast_math: Whether to enable fast math for the render kernel.
     use_shadows: Whether to use shadows.
     use_ambient_lighting: Top-level ambient switch. When False, skips all
                           ambient contributions, including headlight ambient,
@@ -3948,6 +4167,7 @@ def create_render_context(
     cam_res=cam_res_arr,
     cam_id_map=wp.array(active_cam_indices, dtype=int),
     use_textures=use_textures,
+    use_fast_math=use_fast_math,
     use_shadows=use_shadows,
     use_ambient_lighting=use_ambient_lighting,
     background_color=render_util.pack_rgba_to_uint32(

mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py

@@ -22,6 +22,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
 from mujoco.mjx.third_party.mujoco_warp._src.types import OverflowType
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
+wp.set_module_options({"default_grid_stride": False})
+
 
 @wp.kernel
 def _tree_edges(

mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py

@@ -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:
+  # Out:
-  qfrc_spring_out: wp.array2d[float],
+  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:
+  cvel_in: wp.array2d[wp.spatial_vector],
-  qfrc_spring_out: wp.array2d[float],
+  # 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]
+    eid = edgeid - flex_edgeadr[i]
-    if locid >= 0 and locid < flex_edgenum[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))
+  frc = mat43()
-  if flex_bending[bendingadr + 16]:
+  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)
+    ed0 = v1 - v0
-    frc[3] = wp.cross(v1 - v0, v2 - 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))
+  # Gather velocities if damping is enabled
-  for i in range(nvert):
+  vel = mat43()
-    for x in range(3):
+  if not dsbl_damper and flex_damping[f] > 0.0:
-      acc = float(0.0)
+    for j in range(4):
-      for j in range(nvert):
+      bodyid_j = flex_vertbodyid[v[j]]
-        acc += flex_bending[bendingadr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
+      cvel_j = cvel_in[worldid, bodyid_j]
-      force[i, x] = -(acc + flex_bending[bendingadr + 16] * frc[i, x])
+      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):
+  force_spring = mat43()
-    bodyid = flex_vertbodyid[v[i]]
+  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)
+          # Apply forces to body
-          dofnum_i = body_dofnum[bodyid_i]
+          node_pos = flexnode_xpos_in[worldid, nstart + gidx_i]
-          dofadr_i = body_dofadr[bodyid_i]
+          body_xipos = xipos_in[worldid, 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
 
-            if centered or at_origin:
+          offset = body_xipos - node_pos
-              for x in range(3):
+          if not dsbl_spring:
-                if x < dofnum_i:
+            spatial_frc_s = wp.spatial_vector(frc_spring_world, -wp.cross(offset, frc_spring_world))
-                  if not dsbl_spring:
+            wp.atomic_add(flex_spring_body_force_out, worldid, bodyid_i, spatial_frc_s)
-                    wp.atomic_add(
+
-                      qfrc_spring_out,
+          if not dsbl_damper:
-                      worldid,
+            spatial_frc_d = wp.spatial_vector(frc_damper_world, -wp.cross(offset, frc_damper_world))
-                      dofadr_i + x,
+            wp.atomic_add(flex_damper_body_force_out, worldid, bodyid_i, spatial_frc_d)
-                      frc_spring_world[x],
-                    )
-                  if not dsbl_damper:
-                    wp.atomic_add(
-                      qfrc_damper_out,
-                      worldid,
-                      dofadr_i + x,
-                      frc_damper_world[x],
-                    )
 
           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,
+        flex_spring_body_force,
-        d.qfrc_damper,
+        flex_damper_body_force,
         displ_scratch,
         vel_corot_scratch,
       ],
     )
 
+  if m.nflex > 0:
+    if not dsbl_spring:
+      support.apply_ft(m, d, flex_spring_body_force, d.qfrc_spring, True)
+    if not dsbl_damper:
+      support.apply_ft(m, d, flex_damper_body_force, d.qfrc_damper, True)
+
   if m.has_fluid:
     _fluid(m, d)
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py

@@ -26,7 +26,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.func

mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py

@@ -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]
+  u = uv[0] * tex_repeat[0] + offset[0]
-  v = uv[1] * tex_repeat[1]
+  v = uv[1] * tex_repeat[1] + offset[1]
   u = u - wp.floor(u)
   v = v - wp.floor(v)
   tex_color = wp.texture_sample(tex, wp.vec2(u, v), dtype=wp.vec4)
@@ -545,8 +551,6 @@ def render(m: Model, d: Data, rc: RenderContext):
     d: The data on device.
     rc: The render context on device.
   """
-  rc.rgb_data.fill_(rc.background_color)
-  rc.depth_data.fill_(0.0)
   rc.seg_data.fill_(wp.vec2i(-1, -1))
 
   # Specialize the ray-cast helpers to the geom types present in the scene so the
@@ -557,11 +561,11 @@ def render(m: Model, d: Data, rc: RenderContext):
   compute_lighting = _make_compute_lighting(cast_ray_first_hit)
 
   # Static parameters extracted for JAX FFI closure.
-  rc_static = {f.name: getattr(rc, f.name) for f in dataclasses.fields(rc) if f.type in (int, bool, float, wp.vec3)}
+  rc_static = {f.name: getattr(rc, f.name) for f in dataclasses.fields(rc) if f.type in (int, wp.uint32, bool, float, wp.vec3)}
   rc_static["enable_specular_or_emission"] = rc.enable_specular or rc.enable_emission
   M_NLIGHT = m.nlight
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False, module_options={"fast_math": rc.use_fast_math})
   def _render_megakernel(
     # Model:
     geom_type: wp.array[int],
@@ -716,6 +720,8 @@ def render(m: Model, d: Data, rc: RenderContext):
 
     # Early Out
     if geom_id == -1:
+      if render_depth[camid]:
+        depth_out[worldid, depth_adr[camid] + rayid_local] = 0.0
       if wp.static(rc_static["render_skybox"]) and render_rgb[camid]:
         skybox_id = skybox_tex_id[worldid % skybox_tex_id.shape[0]]
         skybox_color = sample_skybox(
@@ -729,6 +735,8 @@ def render(m: Model, d: Data, rc: RenderContext):
           skybox_color[2] * 255.0,
           255.0,
         )
+      elif render_rgb[camid]:
+        rgb_out[worldid, rgb_adr[camid] + rayid_local] = wp.static(rc_static["background_color"])
       return
 
     if render_depth[camid]:

mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py

@@ -19,7 +19,7 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src.types import ProjectionType
 from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.kernel

mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py

@@ -50,7 +50,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.util_misc import poly_potential
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 
 @wp.func
@@ -2140,7 +2140,7 @@ def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3:
   return wp.spatial_bottom(vec) - wp.cross(dif, wp.spatial_top(vec))
 
 
-@wp.kernel
+@wp.kernel(grid_stride=True)
 def _preprocess_tactile_contacts(
   # Model:
   body_weldid: wp.array[int],
@@ -2468,7 +2468,7 @@ def _contact_match(
 
 @cache_kernel
 def _contact_sort(maxmatch: int):
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def contact_sort(
     # Model:
     sensor_intprm: wp.array2d[int],
@@ -2969,7 +2969,7 @@ def energy_pos(m: Model, d: Data):
 
 @cache_kernel
 def _energy_vel_kinetic(nv: int):
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def energy_vel_kinetic(
     # Data in:
     qvel_in: wp.array2d[float],

mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py

@@ -23,7 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import SleepState
 from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
-wp.set_module_options({"enable_backward": False})
+wp.set_module_options({"enable_backward": False, "default_grid_stride": False})
 
 # tree_asleep value for fully awake tree
 K_AWAKE_VAL = -(1 + types.MJ_MINAWAKE)
@@ -244,6 +244,9 @@ def _wake_tree(
   tree_asleep_out: wp.array2d[int],
 ) -> int:
   """Wakes tree treeid and its associated cycle, returning number of woke trees."""
+  if treeid < 0 or treeid >= ntree:
+    return 0
+
   asleep_val = tree_asleep_out[worldid, treeid]
   if asleep_val < 0:
     if wakeval < asleep_val:

mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py

@@ -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(
+    if tile.elemid.size == 0:
-      _tile_cholesky_factorize_block(tile),
+      wp.launch(
-      dim=(d.nworld, tile.adr.size),
+        _small_cholesky_factorize_block(tile.size),
-      inputs=[m.qLD_block_adr, M, tile.elemid, tile.adr],
+        dim=(d.nworld, tile.adr.size),
-      outputs=[L],
+        inputs=[m.M_rowadr, m.qLD_block_adr, M, tile.adr],
-      block_dim=m.block_dim.cholesky_factorize,
+        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),
+  Compact blocks use reciprocal diagonals, full small blocks use scalar Cholesky, larger dense
-  sparse blocks via the LDL factor over the LDL region (offset qLD_block_total), and simple
+  blocks use tile Cholesky, and oversized blocks use sparse LDL.
-  (diagonal) blocks need only D = 1/diag. The passes write disjoint dofs and may all run at once.
   """
-  if m.qLD_has_dense:
+  if m.M_tiles:
-    _factor_block_dense(m, d, d.M, d.qLD)
+    _factor_blocks(m, d, d.M, d.qLD, d.qLDiagInv)
-  if m.qLD_has_sparse:
+  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
+    # Copy y to x_out for sparse-block dofs only.
-    # (diagonal) blocks use the dedicated 1/diag solve.
     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
+@wp.func
-def _solve_simple(
+def _small_cholesky_solve(
   # In:
-  simple_dofs: wp.array[int],
+  block_size: int,
-  D: wp.array2d[float],
+  worldid: int,
-  y: wp.array2d[float],
+  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.
+  for i in range(block_size):
-  worldid, s = wp.tid()
+    value = y_in[worldid, start + i]
-  dofid = simple_dofs[s]
+    for k in range(i):
-  x_out[worldid, dofid] = D[worldid, dofid] * y[worldid, dofid]
+      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
+@cache_kernel
-def _factor_solve_simple(
+def _small_cholesky_solve_block(block_size: int):
-  # Model:
+  @wp.kernel(module="unique", enable_backward=False)
-  M_rownnz: wp.array[int],
+  def kernel(
-  M_rowadr: wp.array[int],
+    # Model:
-  # Data in:
+    qLD_block_adr: wp.array[int],
-  M_in: wp.array2d[float],
+    # In:
-  # In:
+    block_dof: wp.array[int],
-  simple_dofs: wp.array[int],
+    D_in: wp.array2d[float],
-  y: wp.array2d[float],
+    L_in: wp.array2d[float],
-  # Out:
+    y_in: wp.array2d[float],
-  D_out: wp.array2d[float],
+    # Out:
-  x_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, blk = wp.tid()
-  worldid, s = wp.tid()
+    start = block_dof[blk]
-  dofid = simple_dofs[s]
+    size = wp.static(block_size)
-  diag_i = M_rowadr[dofid] + M_rownnz[dofid] - 1
+    factor_adr = qLD_block_adr[start]
-  d_inv = 1.0 / M_in[worldid, diag_i]
+    if factor_adr == Q_LD_BLOCK_COMPACT:
-  D_out[worldid, dofid] = d_inv
+      for i in range(wp.static(block_size)):
-  x_out[worldid, dofid] = d_inv * y[worldid, dofid]
+        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
+    if tile.elemid.size == 0:
-    # for better occupancy while moderate blocks still want a couple warps (16/27->64, 60->32).
+      wp.launch(
-    block_dim = m.block_dim.cholesky_solve if tile.size <= 40 else 32
+        _small_cholesky_solve_block(tile.size),
-    wp.launch_tiled(
+        dim=(d.nworld, tile.adr.size),
-      _tile_cholesky_solve_block(tile),
+        inputs=[m.qLD_block_adr, tile.adr, D, L, y],
-      dim=(d.nworld, tile.adr.size),
+        outputs=[x],
-      inputs=[m.qLD_block_adr, tile.adr, L, y],
+        block_dim=m.block_dim.small_cholesky,
-      outputs=[x],
+      )
-      block_dim=block_dim,
+    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
+  Compact blocks use reciprocal diagonals, full small blocks use scalar Cholesky, dense blocks use
-  blocks from the LDL region (offset qLD_block_total); simple (diagonal) blocks are a plain x = D*y.
+  tile Cholesky, and sparse blocks use the LDL region.
-  The passes write disjoint dofs; the sparse pass skips dense and simple dofs so it does not clobber
-  their results.
 
   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:
+  if m.M_tiles:
-    _solve_block_dense(m, d, L, x, y)
+    _solve_blocks(m, d, L, D, x, y)
-  if m.qLD_has_sparse:
+  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(
+@cache_kernel
-  m: Model, d: Data, M: wp.array2d[float], x: wp.array2d[float], y: wp.array2d[float], L: wp.array2d[float]
+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(
+    if tile.elemid.size == 0:
-      _tile_cholesky_factorize_solve_block(tile),
+      wp.launch(
-      dim=(d.nworld, tile.adr.size),
+        _small_cholesky_factorize_solve_block(tile.size),
-      inputs=[m.qLD_block_adr, M, tile.elemid, tile.adr, y],
+        dim=(d.nworld, tile.adr.size),
-      outputs=[x, L],
+        inputs=[m.M_rowadr, m.qLD_block_adr, M, tile.adr, y],
-      block_dim=m.block_dim.cholesky_factorize_solve,
+        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
+  Compact blocks use reciprocal diagonals, full small blocks use scalar Cholesky, dense blocks use
-  blocks via the LDL region, simple (diagonal) blocks via D = 1/diag. Factorizes M, solves for x.
+  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
+  if m.M_tiles:
-  # (offset qLD_block_total); simple blocks via 1/diag. The passes write disjoint dofs.
+    _factor_solve_blocks(m, d, M, L, D, x, y)
-  if m.qLD_has_dense:
+  if L.shape[1] > m.qLD_block_total:
-    _factor_solve_block_dense(m, d, M, x, y, L)
-  if m.qLD_has_sparse:
     L_ldl = L[:, m.qLD_block_total :]
     _factor_i_sparse(m, d, M, L_ldl, D)
     _solve_LD_sparse(m, d, L_ldl, D, x, y)
-  if m.qLD_has_simple:
-    wp.launch(
-      _factor_solve_simple,
-      dim=(d.nworld, m.qLD_simple_dofs.size),
-      inputs=[m.M_rownnz, m.M_rowadr, M, m.qLD_simple_dofs, y],
-      outputs=[D, x],
-    )
 
 
 @cache_kernel

mjx/mujoco/mjx/third_party/mujoco_warp/_src/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_x = float(-1) if li == 0 else float(1)
+          dphi_y = float(-1) if lj == 0 else float(1)
-            dphi_y = float(-1) if lj == 0 else float(1)
+          dphi_z = float(-1) if lk == 0 else float(1)
-            dphi_z = float(-1) if lk == 0 else float(1)
+          phi_x = float(0.5)
-            phi_x = float(0.5)
+          phi_y = float(0.5)
-            phi_y = float(0.5)
+          phi_z = 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
 
           grad_x = dphi_x * phi_y * phi_z
           grad_y = phi_x * dphi_y * phi_z
@@ -175,7 +152,7 @@ def compute_interp_cell_quat(
 
 @cache_kernel
 def mul_m_kernel(check_skip: bool):
-  @wp.kernel(module="unique")
+  @wp.kernel(module="unique", grid_stride=False)
   def _mul_m(
     # Model:
     M_mulm_rowadr: wp.array[int],
@@ -212,7 +189,7 @@ def mul_m_kernel(check_skip: bool):
 
 @cache_kernel
 def mul_m_dense(nv: int, check_skip: bool):
-  @wp.kernel(module="unique")
+  @wp.kernel(module="unique", grid_stride=False)
   def _mul_m_dense(
     # Data in:
     M_in: wp.array3d[float],  # kernel_analyzer: ignore
@@ -550,7 +527,7 @@ def jac_dof(
 
 @cache_kernel
 def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def _jac(
     # Model:
     body_parentid: wp.array[int],
@@ -702,7 +679,7 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
   if sig >= (1 << State.NSTATE):
     raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def _get_state(
     # Model:
     nq: int,
@@ -857,7 +834,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
   if sig >= (1 << State.NSTATE):
     raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, grid_stride=False)
   def _set_state(
     # Model:
     nq: int,
@@ -1083,3 +1060,221 @@ def select_top4_weights(
     )
 
   return selected_b, selected_W
+
+
+@wp.func
+def get_face_metadata(
+  # In:
+  cellnum_x: int,
+  cellnum_y: int,
+  cellnum_z: int,
+  face_elem_idx: int,
+  order_abs: int,
+) -> Tuple[int, int, int, int, int, int]:
+  size01 = cellnum_y * cellnum_z
+  size23 = cellnum_x * cellnum_z
+  size45 = cellnum_x * cellnum_y
+
+  face_id = 0
+  within_face = 0
+
+  if face_elem_idx < size01:
+    face_id = 0
+    within_face = face_elem_idx
+  elif face_elem_idx < 2 * size01:
+    face_id = 1
+    within_face = face_elem_idx - size01
+  elif face_elem_idx < 2 * size01 + size23:
+    face_id = 2
+    within_face = face_elem_idx - 2 * size01
+  elif face_elem_idx < 2 * size01 + 2 * size23:
+    face_id = 3
+    within_face = face_elem_idx - 2 * size01 - size23
+  elif face_elem_idx < 2 * size01 + 2 * size23 + size45:
+    face_id = 4
+    within_face = face_elem_idx - 2 * size01 - 2 * size23
+  else:
+    face_id = 5
+    within_face = face_elem_idx - 2 * size01 - 2 * size23 - size45
+
+  normal_axis = face_id // 2
+
+  c1 = 0
+  if face_id == 0 or face_id == 1:
+    c1 = cellnum_z
+  elif face_id == 2 or face_id == 3:
+    c1 = cellnum_x
+  else:
+    c1 = cellnum_y
+
+  fixed_dim = 0
+  if normal_axis == 0:
+    fixed_dim = cellnum_x
+  elif normal_axis == 1:
+    fixed_dim = cellnum_y
+  else:
+    fixed_dim = cellnum_z
+  g_fixed = (face_id % 2) * fixed_dim * order_abs
+
+  q0 = within_face // c1
+  q1 = within_face % c1
+
+  ny_g = cellnum_y * order_abs + 1
+  nz_g = cellnum_z * order_abs + 1
+
+  return normal_axis, g_fixed, q0, q1, ny_g, nz_g
+
+
+@wp.func
+def gather_face_node_index_fast(
+  # In:
+  normal_axis: int,
+  g_fixed: int,
+  q0: int,
+  q1: int,
+  ny_g: int,
+  nz_g: int,
+  local_idx: int,
+  order_abs: int,
+) -> int:
+  l0 = local_idx // (order_abs + 1)
+  l1 = local_idx % (order_abs + 1)
+
+  g = wp.vec3i(0, 0, 0)
+  if normal_axis == 0:
+    g = wp.vec3i(g_fixed, q0 * order_abs + l0, q1 * order_abs + l1)
+  elif normal_axis == 1:
+    g = wp.vec3i(q1 * order_abs + l1, g_fixed, q0 * order_abs + l0)
+  else:
+    g = wp.vec3i(q0 * order_abs + l0, q1 * order_abs + l1, g_fixed)
+
+  return g[0] * ny_g * nz_g + g[1] * nz_g + g[2]
+
+
+@wp.func
+def gather_face_node_index(
+  # In:
+  cellnum_x: int,
+  cellnum_y: int,
+  cellnum_z: int,
+  face_elem_idx: int,
+  local_idx: int,
+  order_abs: int,
+) -> int:
+  normal_axis, g_fixed, q0, q1, ny_g, nz_g = get_face_metadata(cellnum_x, cellnum_y, cellnum_z, face_elem_idx, order_abs)
+  return gather_face_node_index_fast(normal_axis, g_fixed, q0, q1, ny_g, nz_g, local_idx, order_abs)
+
+
+@wp.func
+def compute_interp_face_quat(
+  # Data in:
+  flexnode_xpos_in: wp.array2d[wp.vec3],
+  # In:
+  cellnum_x: int,
+  cellnum_y: int,
+  cellnum_z: int,
+  face_elem_idx: int,
+  nstart: int,
+  order_abs: int,
+  worldid: int,
+) -> wp.quat:
+  normal_axis, g_fixed, q0, q1, ny_g, nz_g = get_face_metadata(
+    cellnum_x,
+    cellnum_y,
+    cellnum_z,
+    face_elem_idx,
+    order_abs,
+  )
+
+  t1 = wp.vec3(0.0)
+  t2 = wp.vec3(0.0)
+
+  npc = (order_abs + 1) * (order_abs + 1)
+
+  for local_idx in range(9):
+    if local_idx < npc:
+      gidx = gather_face_node_index_fast(
+        normal_axis,
+        g_fixed,
+        q0,
+        q1,
+        ny_g,
+        nz_g,
+        local_idx,
+        order_abs,
+      )
+      node_pos = flexnode_xpos_in[worldid, nstart + gidx]
+
+      l0 = local_idx // (order_abs + 1)
+      l1 = local_idx % (order_abs + 1)
+
+      dphi0 = -1.0 + 2.0 * float(l0)
+      dphi1 = -1.0 + 2.0 * float(l1)
+      phi0 = 0.5
+      phi1 = 0.5
+
+      grad0 = dphi0 * phi1
+      grad1 = phi0 * dphi1
+
+      t1 += node_pos * grad0
+      t2 += node_pos * grad1
+
+  normal = wp.cross(t1, t2)
+
+  F = wp.mat33(0.0)
+  if normal_axis == 0:
+    F = wp.mat33(
+      normal[0],
+      t1[0],
+      t2[0],
+      normal[1],
+      t1[1],
+      t2[1],
+      normal[2],
+      t1[2],
+      t2[2],
+    )
+  elif normal_axis == 1:
+    F = wp.mat33(
+      t2[0],
+      normal[0],
+      t1[0],
+      t2[1],
+      normal[1],
+      t1[1],
+      t2[2],
+      normal[2],
+      t1[2],
+    )
+  else:
+    F = wp.mat33(
+      t1[0],
+      t2[0],
+      normal[0],
+      t1[1],
+      t2[1],
+      normal[1],
+      t1[2],
+      t2[2],
+      normal[2],
+    )
+
+  return mat33_to_quat_polar(F)
+
+
+@wp.func
+def flex_phi(s: float, i: int, order: int) -> float:
+  return 1.0 - s if i == 0 else s
+
+
+@wp.func
+def flex_dphi(s: float, i: int, order: int) -> float:
+  return -1.0 if i == 0 else 1.0
+
+
+@wp.func
+def dphi2D(s0: float, l0: int, s1: float, l1: int, order: int, direction: int) -> float:
+  if direction == 0:
+    return flex_dphi(s0, l0, order) * flex_phi(s1, l1, order)
+  else:
+    return flex_phi(s0, l0, order) * flex_dphi(s1, l1, order)

mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_flex.py

@@ -1,351 +0,0 @@
-# Copyright 2025 The Newton Developers
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Tests for flex."""
-
-import mujoco
-import numpy as np
-from absl.testing import absltest
-from absl.testing import parameterized
-
-import mujoco.mjx.third_party.mujoco_warp as mjw
-from mujoco.mjx.third_party.mujoco_warp import test_data
-
-_TRILINEAR_STRAIN_XML = """
-<mujoco>
-  <option gravity="0 0 -9.81">
-    <flag contact="disable"/>
-  </option>
-  <worldbody>
-    <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
-              pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
-              dof="trilinear">
-      <edge equality="strain"/>
-      <contact selfcollide="none"/>
-    </flexcomp>
-  </worldbody>
-</mujoco>
-"""
-
-# tolerance for difference between MuJoCo and MJWarp, mostly due to float32
-_TOLERANCE = 5e-4
-
-
-class TrilinearFlexTest(parameterized.TestCase):
-  def test_flexstrain_constraint_at_rest(self):
-    """Test FLEXSTRAIN constraint count and residuals match MuJoCo at rest."""
-    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
-    mjd = mujoco.MjData(mjm)
-    mujoco.mj_forward(mjm, mjd)
-
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd)
-
-    mjw.fwd_position(m, d)
-    mjw.make_constraint(m, d)
-
-    # constraint counts should match
-    ne_warp = d.ne.numpy()[0]
-    self.assertEqual(ne_warp, mjd.ne, f"ne mismatch: warp={ne_warp}, mj={mjd.ne}")
-
-    nefc_warp = d.nefc.numpy()[0]
-    self.assertEqual(nefc_warp, mjd.nefc, f"nefc mismatch: warp={nefc_warp}, mj={mjd.nefc}")
-
-    # residuals should match
-    efc_pos = d.efc.pos.numpy()[0, :nefc_warp]
-    efc_pos_mj = mjd.efc_pos[: mjd.nefc]
-    np.testing.assert_allclose(efc_pos, efc_pos_mj, atol=1e-5, err_msg="FLEXSTRAIN residuals should match MuJoCo at rest")
-
-  def test_flexstrain_constraint_perturbed(self):
-    """Test FLEXSTRAIN residuals and Jacobians match MuJoCo under perturbation."""
-    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
-    mjd = mujoco.MjData(mjm)
-
-    # perturb first node
-    mjd.qpos[0] += 0.01
-    mjd.qpos[3] += 0.005
-    mujoco.mj_forward(mjm, mjd)
-
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd)
-
-    mjw.fwd_position(m, d)
-    mjw.make_constraint(m, d)
-
-    nefc = d.nefc.numpy()[0]
-    nv = mjm.nv
-    self.assertEqual(nefc, mjd.nefc)
-
-    # residuals
-    efc_pos_warp = d.efc.pos.numpy()[0, :nefc]
-    efc_pos_mj = mjd.efc_pos[: mjd.nefc]
-    np.testing.assert_allclose(
-      efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under perturbation"
-    )
-
-    # Jacobians
-    if mujoco.mj_isSparse(mjm):
-      mj_efc_J = np.zeros((mjd.nefc, nv))
-      mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind)
-    else:
-      mj_efc_J = mjd.efc_J.reshape((mjd.nefc, nv))
-
-    if m.is_sparse:
-      warp_efc_J = np.zeros((nefc, nv))
-      mujoco.mju_sparse2dense(
-        warp_efc_J,
-        d.efc.J.numpy()[0, 0],
-        d.efc.J_rownnz.numpy()[0, :nefc],
-        d.efc.J_rowadr.numpy()[0, :nefc],
-        d.efc.J_colind.numpy()[0, 0],
-      )
-    else:
-      warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv]
-
-    np.testing.assert_allclose(warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo")
-
-  def test_flexstrain_constraint_rotated(self):
-    """Test FLEXSTRAIN residuals and Jacobians match MuJoCo under large rotation perturbation."""
-    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
-    mjd = mujoco.MjData(mjm)
-
-    # Apply a rotation perturbation: rotate all node positions around Y axis by 30 degrees
-    # (0.5235 radians)
-    theta = 0.5235
-    cos_t = np.cos(theta)
-    sin_t = np.sin(theta)
-    for i in range(0, mjm.nq, 3):
-      x = mjd.qpos[i]
-      z = mjd.qpos[i + 2]
-      mjd.qpos[i] = x * cos_t - z * sin_t
-      mjd.qpos[i + 2] = x * sin_t + z * cos_t
-
-    mujoco.mj_forward(mjm, mjd)
-
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd)
-
-    mjw.fwd_position(m, d)
-    mjw.make_constraint(m, d)
-
-    nefc = d.nefc.numpy()[0]
-    nv = mjm.nv
-    self.assertEqual(nefc, mjd.nefc)
-
-    # residuals
-    efc_pos_warp = d.efc.pos.numpy()[0, :nefc]
-    efc_pos_mj = mjd.efc_pos[: mjd.nefc]
-    np.testing.assert_allclose(
-      efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under rotation"
-    )
-
-    # Jacobians
-    if mujoco.mj_isSparse(mjm):
-      mj_efc_J = np.zeros((mjd.nefc, nv))
-      mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind)
-    else:
-      mj_efc_J = mjd.efc_J.reshape((mjd.nefc, nv))
-
-    if m.is_sparse:
-      warp_efc_J = np.zeros((nefc, nv))
-      mujoco.mju_sparse2dense(
-        warp_efc_J,
-        d.efc.J.numpy()[0, 0],
-        d.efc.J_rownnz.numpy()[0, :nefc],
-        d.efc.J_rowadr.numpy()[0, :nefc],
-        d.efc.J_colind.numpy()[0, 0],
-      )
-    else:
-      warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv]
-
-    np.testing.assert_allclose(
-      warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo under rotation"
-    )
-
-  def test_flexstrain_rotational_invariance(self):
-    """Test that FLEXSTRAIN residuals are invariant under rigid translation."""
-    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
-    mjd = mujoco.MjData(mjm)
-    mujoco.mj_forward(mjm, mjd)
-
-    # Get reference residuals
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd)
-    mjw.fwd_position(m, d)
-    mjw.make_constraint(m, d)
-    nefc = d.nefc.numpy()[0]
-    efc_pos_rest = d.efc.pos.numpy()[0, :nefc].copy()
-
-    # Apply uniform translation to all nodes (rigid motion)
-    mjd2 = mujoco.MjData(mjm)
-    # All flex nodes have 3 DOFs (slide joints), shift all x by 0.1
-    for i in range(0, mjm.nq, 3):
-      mjd2.qpos[i] += 0.1  # shift x
-    mujoco.mj_forward(mjm, mjd2)
-
-    d2 = mjw.put_data(mjm, mjd2)
-    mjw.fwd_position(m, d2)
-    mjw.make_constraint(m, d2)
-    nefc2 = d2.nefc.numpy()[0]
-    efc_pos_shifted = d2.efc.pos.numpy()[0, :nefc2]
-
-    # Residuals should remain near zero for rigid translation
-    np.testing.assert_allclose(
-      efc_pos_shifted, efc_pos_rest, atol=1e-4, err_msg="FLEXSTRAIN residuals should be invariant under rigid translation"
-    )
-
-  def test_trilinear_gravity_parity(self):
-    """Test that trilinear simulation matches MuJoCo after multiple steps."""
-    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
-
-    # MuJoCo reference
-    mjd = mujoco.MjData(mjm)
-    for _ in range(10):
-      mujoco.mj_step(mjm, mjd)
-
-    # Warp
-    mjd_warp = mujoco.MjData(mjm)
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd_warp)
-    z0 = d.qpos.numpy()[0, 2]
-    for _ in range(10):
-      mjw.step(m, d)
-
-    qpos_warp = d.qpos.numpy()[0]
-    qpos_mj = mjd.qpos
-
-    # The cube should have fallen
-    self.assertLess(qpos_warp[2], z0, "Cube should fall under gravity")
-
-    # Allow larger tolerance for accumulated integration error
-    np.testing.assert_allclose(qpos_warp, qpos_mj, atol=0.01, err_msg="Trilinear qpos diverges from MuJoCo after 10 steps")
-
-  def test_trilinear_node_positions(self):
-    """Test that flexnode_xpos are computed correctly from body kinematics."""
-    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
-    mjd = mujoco.MjData(mjm)
-    mujoco.mj_forward(mjm, mjd)
-
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd)
-    mjw.fwd_position(m, d)
-
-    # Compute expected node positions: xpos_n = body_xpos + body_xmat @ flex_node
-    nflexnode = mjm.nflexnode
-    warp_xpos = d.flexnode_xpos.numpy()[0, :nflexnode]
-    nodeadr = mjm.flex_nodeadr[0]
-    nodenum = mjm.flex_nodenum[0]
-    for n in range(nodenum):
-      bodyid = mjm.flex_nodebodyid[nodeadr + n]
-      body_xpos = mjd.xpos[bodyid]
-      body_xmat = mjd.xmat[bodyid].reshape(3, 3)
-      node_local = mjm.flex_node[nodeadr + n]
-      expected = body_xpos + body_xmat @ node_local
-      np.testing.assert_allclose(warp_xpos[n], expected, atol=1e-5, err_msg=f"flexnode_xpos mismatch for node {n}")
-
-  def test_trilinear_passive_forces_parity(self):
-    """Test passive forces (elasticity) match MuJoCo for trilinear flex."""
-    xml = """
-    <mujoco>
-      <worldbody>
-        <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
-                  pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
-                  dof="trilinear">
-          <elasticity young="1e4" poisson="0.1" damping="0.01"/>
-          <contact selfcollide="none"/>
-        </flexcomp>
-      </worldbody>
-    </mujoco>
-    """
-    mjm = mujoco.MjModel.from_xml_string(xml)
-    mjd = mujoco.MjData(mjm)
-
-    # perturb first node to generate non-zero elasticity forces
-    mjd.qpos[0] += 0.01
-    mjd.qpos[3] += 0.005
-    mujoco.mj_forward(mjm, mjd)
-
-    m = mjw.put_model(mjm)
-    d = mjw.put_data(mjm, mjd)
-
-    mjw.fwd_position(m, d)
-    mjw.passive(m, d)
-
-    qfrc_passive_warp = d.qfrc_passive.numpy()[0]
-    qfrc_passive_mj = mjd.qfrc_passive
-
-    # Verify they match
-    np.testing.assert_allclose(
-      qfrc_passive_warp, qfrc_passive_mj, atol=_TOLERANCE, err_msg="qfrc_passive mismatch for trilinear flex with elasticity"
-    )
-
-  @parameterized.parameters("strain", "true")
-  def test_trilinear_equality_types(self, equality):
-    """Test trilinear with different equality types."""
-    xml = f"""
-    <mujoco>
-      <option gravity="0 0 -9.81">
-        <flag contact="disable"/>
-      </option>
-      <worldbody>
-        <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
-                  pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
-                  dof="trilinear">
-          <edge equality="{equality}"/>
-          <contact selfcollide="none"/>
-        </flexcomp>
-      </worldbody>
-    </mujoco>
-    """
-    mjm, mjd, m, d = test_data.fixture(xml=xml)
-
-    # Should not crash
-    mjw.forward(m, d)
-
-    # Constraint count should match
-    self.assertEqual(d.nefc.numpy()[0], mjd.nefc)
-
-  def test_trilinear_contact_qfrc_constraint(self):
-    """Test qfrc_constraint parity for trilinear flex with ground contacts."""
-    xml = """
-    <mujoco>
-      <option gravity="0 0 -9.81"/>
-      <worldbody>
-        <geom type="plane" size="1 1 0.1"/>
-        <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
-                  pos="0 0 0.05" name="cube" dim="3" mass="1" radius="0.02"
-                  dof="trilinear">
-          <edge equality="strain"/>
-          <contact selfcollide="none"/>
-        </flexcomp>
-      </worldbody>
-    </mujoco>
-    """
-    mjm, mjd, m, d = test_data.fixture(xml=xml)
-    mjw.forward(m, d)
-
-    # Verify contacts are generated
-    nacon = d.nacon.numpy()[0]
-    self.assertGreater(nacon, 0, "Expected contacts between flex and plane")
-    self.assertEqual(nacon, mjd.ncon)
-
-    # Verify qfrc_constraint parity
-    qfrc_warp = d.qfrc_constraint.numpy()[0]
-    qfrc_mj = mjd.qfrc_constraint
-    np.testing.assert_allclose(qfrc_warp, qfrc_mj, atol=1e-4, err_msg="qfrc_constraint mismatch for trilinear flex contacts")
-
-
-if __name__ == "__main__":
-  absltest.main()

mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py

@@ -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
+      factorization is a per-block decision -- see M_tiles 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)
     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_block_adr: packed factor offset; Q_LD_BLOCK_* sentinel otherwise (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,)
     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
+    M_tiles: scalar and tiled block-factorization groups
-    qLD_updates: tuple of index triples for sparse factorization
+    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_dynprm: array("*", "nu", vec10)
-  actuator_gainprm: array("*", "nu", vec10f)
+  actuator_gainprm: array("*", "nu", vec10)
-  actuator_biasprm: array("*", "nu", vec10f)
+  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_block_adr: array("nv", int)
-  qLD_dof_dense: wp.array[int]
-  qLD_dof_simple: wp.array[int]
-  qLD_simple_dofs: wp.array[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_tree: tuple[array("nbody", int), ...]
-  body_branches: wp.array[int]
+  body_branches: array("nbody_branches", int)
-  body_branch_start: wp.array[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_ellipsoid_adr: array("nbody_fluid_ellipsoid", int)
-  body_fluid_box_adr: wp.array[int]
+  body_fluid_box_adr: array("nbody_fluid_box", int)
-  jnt_limited_slide_hinge_adr: wp.array[int]
+  jnt_limited_slide_hinge_adr: array("njnt_limited_slide_hinge", int)
-  jnt_limited_ball_adr: wp.array[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_row: array("ndof_tri", int)
-  dof_tri_col: wp.array[int]
+  dof_tri_col: array("ndof_tri", int)
-  nxn_geom_pair: wp.array[wp.vec2i]
+  nxn_geom_pair: array("nnxn_geom_pair", wp.vec2i)
-  nxn_geom_pair_filtered: wp.array[wp.vec2i]
+  nxn_geom_pair_filtered: array("nnxn_geom_pair_filtered", wp.vec2i)
-  nxn_pairid: wp.array[wp.vec2i]
+  nxn_pairid: array("nnxn_geom_pair", wp.vec2i)
-  nxn_pairid_filtered: wp.array[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_connect_adr: array("neq_connect", int)
-  eq_wld_adr: wp.array[int]
+  eq_wld_adr: array("neq_wld", int)
-  eq_jnt_adr: wp.array[int]
+  eq_jnt_adr: array("neq_jnt", int)
-  eq_ten_adr: wp.array[int]
+  eq_ten_adr: array("neq_ten", int)
-  eq_flex_adr: wp.array[int]
+  eq_flex_adr: array("neq_flex", int)
-  eq_flexstrain_adr: wp.array[int]
+  eq_flexstrain_adr: array("neq_flexstrain", int)
-  tendon_jnt_adr: wp.array[int]
+  tendon_jnt_adr: array("ntendon_jnt", int)
-  tendon_site_pair_adr: wp.array[int]
+  tendon_site_pair_adr: array("ntendon_site_pair", int)
-  tendon_geom_adr: wp.array[int]
+  tendon_geom_adr: array("ntendon_geom", int)
-  tendon_limited_adr: wp.array[int]
+  tendon_limited_adr: array("ntendon_limited", int)
   max_ten_J_rownnz: int
-  jtcj_max_pairs: int
+  ten_wrapadr_site: array("nten_wrapadr_site", int)
-  ten_wrapadr_site: wp.array[int]
+  ten_wrapnum_site: array("ntendon", int)
-  ten_wrapnum_site: wp.array[int]
+  wrap_jnt_adr: array("nwrap_jnt", int)
-  wrap_jnt_adr: wp.array[int]
+  wrap_site_adr: array("nwrap_site", int)
-  wrap_site_adr: wp.array[int]
+  wrap_site_pair_adr: array("nwrap_site_pair", int)
-  wrap_site_pair_adr: wp.array[int]
+  wrap_geom_adr: array("nwrap_geom", int)
-  wrap_geom_adr: wp.array[int]
   wrap_pulley_scale: array("nwrap", float)
-  actuator_trntype_body_adr: wp.array[int]
+  actuator_trntype_body_adr: array("nacttrnbody", int)
-  sensor_pos_adr: wp.array[int]
+  sensor_pos_adr: array("nsensor_pos", int)
-  sensor_limitpos_adr: wp.array[int]
+  sensor_limitpos_adr: array("nsensor_limitpos", int)
-  sensor_vel_adr: wp.array[int]
+  sensor_vel_adr: array("nsensor_vel", int)
-  sensor_limitvel_adr: wp.array[int]
+  sensor_limitvel_adr: array("nsensor_limitvel", int)
-  sensor_acc_adr: wp.array[int]
+  sensor_acc_adr: array("nsensor_acc", int)
-  sensor_rangefinder_adr: wp.array[int]
+  sensor_rangefinder_adr: array("nrangefinder", int)
-  rangefinder_sensor_adr: wp.array[int]
+  rangefinder_sensor_adr: array("nsensor", int)
-  sensor_collision_start_adr: wp.array[int]
+  sensor_collision_start_adr: array("nsensor_collision_start_adr", int)
   collision_sensor_adr: array("nsensor", int)
-  sensor_touch_adr: wp.array[int]
+  sensor_touch_adr: array("nsensor_touch", int)
-  sensor_limitfrc_adr: wp.array[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_updates: tuple[array("nqLD_all_updates", wp.vec3i), ...]
-  qLD_all_updates: wp.array[wp.vec3i]
+  qLD_all_updates: array("nqLD_all_updates", wp.vec3i)
-  qLD_level_offsets: wp.array[int]
+  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_i: array("nM_fullm", int)
-  M_fullm_j: wp.array[int]
+  M_fullm_j: array("nM_fullm", int)
-  M_elemid: wp.array2d[int]  # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
+  M_elemid: array("nv", "nv", 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_hinit_i: array("nC", 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_i: array("nM_fullm_upper", int)
-  M_fullm_upper_j: wp.array[int]
+  M_fullm_upper_j: array("nM_fullm_upper", int)
-  M_fullm_upper_elemid: wp.array[int]
+  M_fullm_upper_elemid: array("nM_fullm_upper", int)
-  qD_fullm_i: wp.array[int]  # D-structure (full square) row indices for RNE derivatives
+  qD_fullm_i: array("nqD_fullm", 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
+  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_rowadr: array("nv_plus_1", int)  # start address for each row [nv+1]
-  M_mulm_col: wp.array[int]  # column index to gather from
+  M_mulm_col: array("nM_mulm", int)  # column index to gather from
-  M_mulm_madr: wp.array[int]  # matrix address to read
+  M_mulm_madr: array("nM_mulm", int)  # matrix address to read
-  flexelem_geom_pair_filtered: wp.array[wp.vec2i]
+  flexelem_geom_pair_filtered: array("nflexelem_geom_pair_filtered", wp.vec2i)
-  flexvert_geom_pair_filtered: wp.array[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_adr: first efc row of each JTDAJ block      (nworld, njmax)
-    jtdaj_nrow: efc rows per JTDAJ block            (nworld, njmax)
+    jtdaj_nrow: efc rows per JTDAJ block              (nworld, njmax)
-    jtdaj_nblock: number of JTDAJ blocks             (nworld,)
+    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_colind: array("nworld", 1, "njmax_nnz", int)
-  J: wp.array3d[float]
+  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_row: row index of compacted Hessian dof-pairs      (nvmax_pad_sq,)
-    cdof_tri_col: col index of compacted Hessian dof-pairs      (nvmax_pad^2,)
+    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]
+  M: array("nworld", "nC", float)
-  qLD: wp.array2d[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]
+  ctol: array(1, float)
-  cls_tol: wp.array[float]
+  cls_tol: array(1, float)
-  cdof_tri_row: wp.array[int]
+  cdof_tri_row: array("nvmax_pad_sq", int)
-  cdof_tri_col: wp.array[int]
+  cdof_tri_col: array("nvmax_pad_sq", int)
-  cM: wp.array3d[float]
+  cM: array("nworld", "nvmax_pad", "nvmax_pad", float)
-  cqLD: wp.array3d[float]
+  cqLD: array("nworld", "nvmax_pad", "nvmax_pad", float)
-  crhs: wp.array3d[float]
+  crhs: array("nworld", "nvmax_pad", 1, float)
-  cx: wp.array3d[float]
+  cx: array("nworld", "nvmax_pad", 1, float)
-  cJ: wp.array3d[float]
+  cJ: array("nworld", "njmax_pad", "nvmax_pad", float)
-  cMa: wp.array2d[float]
+  cMa: array("nworld", "nvmax_pad", float)
-  cqfrc_smooth: wp.array2d[float]
+  cqfrc_smooth: array("nworld", "nvmax_pad", float)
-  cqacc_smooth: wp.array2d[float]
+  cqacc_smooth: array("nworld", "nvmax_pad", float)
-  cqacc_warmstart: wp.array2d[float]
+  cqacc_warmstart: array("nworld", "nvmax_pad", float)
-  cqacc: wp.array2d[float]
+  cqacc: array("nworld", "nvmax_pad", float)
-  cqfrc_constraint: wp.array2d[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]
+  quad_changed_ids: wp.array2d[int]
-  changed_efc_count: wp.array[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)
+  quad_changed_ids: wp.array2d[int]
-  changed_efc_ids: wp.array2d[int]
+  quad_changed_count: wp.array[int]
-  changed_efc_count: wp.array[int]
+  # the full-coordinate Data, set by solve_compact (None natively)
+  compact_m_full: Optional["Model"] = None
+  compact_d_full: Optional["Data"] = None
 
 
 @dataclasses.dataclass
@@ -2373,6 +2409,7 @@ class RenderContext:
     cam_res: camera resolution for actively rendering cameras
     cam_id_map: camera id map
     use_textures: whether to use textures
+    use_fast_math: whether to enable fast math for the render kernel
     use_shadows: whether to use shadows
     use_ambient_lighting: top-level switch for ambient contributions
     background_color: color used for missed rays when no skybox is rendered
@@ -2460,6 +2497,7 @@ class RenderContext:
   cam_res: array("ncam", wp.vec2i)
   cam_id_map: array("ncam", int)
   use_textures: bool
+  use_fast_math: bool
   use_shadows: bool
   use_ambient_lighting: bool
   background_color: wp.uint32

mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml

@@ -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",
+    "mujoco>=3.11.0",
-    "warp-lang>=1.14",
+    "warp-lang>=1.15",
     "mjviser>=0.0.10",
     "pillow",
 ]

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))

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,

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,
   )

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,

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,
   )

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
+    elemid: flat CSR gather indices for tiled blocks; empty selects the
-      (absent -> nC sentinel)
+      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,
+        'cJ': True,
-        'cM': False,
+        'cM': True,
-        'cMa': False,
+        '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,
+        'cqLD': True,
-        'cqacc': False,
+        'cqacc': True,
-        'cqacc_smooth': False,
+        'cqacc_smooth': True,
-        'cqacc_warmstart': False,
+        'cqacc_warmstart': True,
-        'cqfrc_constraint': False,
+        'cqfrc_constraint': True,
-        'cqfrc_smooth': False,
+        'cqfrc_smooth': True,
         'crb': True,
-        'crhs': False,
+        'crhs': True,
         'ctol': False,
         'ctrl': True,
         'cvel': True,
-        'cx': False,
+        'cx': True,
         'dof_awake_ind': True,
         'dof_cdof': True,
         'dof_island': True,
@@ -1394,6 +1391,8 @@ _BATCH_DIM = {
         'efc_islandid': True,
         'energy': True,
         'eq_active': True,
+        'face_quat': True,
+        'face_xpos': True,
         'flex_aabb_max': True,
         'flex_aabb_min': True,
         'flexedge_J': True,
@@ -1547,6 +1546,7 @@ _BATCH_DIM = {
         'block_dim__ray': False,
         'block_dim__render': False,
         'block_dim__segmented_sort': False,
+        'block_dim__small_cholesky': False,
         'block_dim__solve_LD_sparse_fused': False,
         'block_dim__solve_beta_accumulate': False,
         'block_dim__solve_init_search_cg': False,
@@ -1581,6 +1581,7 @@ _BATCH_DIM = {
         'body_pos': True,
         'body_quat': True,
         'body_rootid': False,
+        'body_simple': False,
         'body_subtreemass': True,
         'body_tree': False,
         'body_treeid': False,
@@ -1629,6 +1630,7 @@ _BATCH_DIM = {
         'eq_type': False,
         'eq_wld_adr': False,
         'exclude_signature': False,
+        'flex_bend_interp_map': False,
         'flex_bending': False,
         'flex_bendingadr': False,
         'flex_cell_map': False,
@@ -1655,6 +1657,9 @@ _BATCH_DIM = {
         'flex_evpairadr': False,
         'flex_evpairflexid': False,
         'flex_evpairnum': False,
+        'flex_face': False,
+        'flex_face_map': False,
+        'flex_faceadr': False,
         'flex_friction': False,
         'flex_gap': False,
         'flex_internal': False,
@@ -1747,7 +1752,6 @@ _BATCH_DIM = {
         'jnt_stiffness': True,
         'jnt_stiffnesspoly': True,
         'jnt_type': False,
-        'jtcj_max_pairs': False,
         'light_active': True,
         'light_ambient': True,
         'light_attenuation': True,
@@ -1815,12 +1819,14 @@ _BATCH_DIM = {
         'neq_flexstrain': False,
         'nexclude': False,
         'nflex': False,
+        'nflexbend_interp': False,
         'nflexbending': False,
         'nflexedge': False,
         'nflexelem': False,
         'nflexelemdata': False,
         'nflexelemedge': False,
         'nflexevpair': False,
+        'nflexface': False,
         'nflexintcell': False,
         'nflexnode': False,
         'nflexshelldata': False,
@@ -1912,13 +1918,7 @@ _BATCH_DIM = {
         'qLD_all_updates': False,
         'qLD_block_adr': False,
         'qLD_block_total': False,
-        'qLD_dof_dense': False,
-        'qLD_dof_simple': False,
-        'qLD_has_dense': False,
-        'qLD_has_simple': False,
-        'qLD_has_sparse': False,
         'qLD_level_offsets': False,
-        'qLD_simple_dofs': False,
         'qLD_updates': False,
         'qpos0': True,
         'qpos_spring': True,

mjx/mujoco/mjx/warp/types_test.py

@@ -14,11 +14,80 @@
 # ==============================================================================
 """Tests for generated MJX Warp types."""
 
+import dataclasses
 from absl.testing import absltest
 from absl.testing import parameterized
+import mujoco
+import mujoco.mjx.warp as mjxw
+from mujoco.mjx.warp import mjwp_io
+from mujoco.mjx.warp import types
+from mujoco.mjx.warp import warp as wp
 import numpy as np
 
-from mujoco.mjx.warp import types
+
+class GeneratedTypesMetadataTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if not mjxw.WARP_INSTALLED or mjwp_io is None:
+      self.skipTest('warp is not installed')
+
+  def test_batched_fields_metadata(self):
+    m = mujoco.MjModel.from_xml_string('<mujoco/>')
+    d = mujoco.MjData(m)
+    mw = mjwp_io.put_model(m)
+    dw = mjwp_io.put_data(m, d)
+
+    def _check_fields(cls_name, obj, warp_cls, prefix=''):
+      for f in dataclasses.fields(obj):
+        val = getattr(obj, f.name)
+        if dataclasses.is_dataclass(val):
+          sub_cls_name = type(val).__name__
+          if sub_cls_name in types._BATCH_DIM and hasattr(
+              types, f'{sub_cls_name}Warp'
+          ):
+            # Dedicated top-level sub-class (Option, Statistic)
+            _check_fields(
+                sub_cls_name, val, getattr(types, f'{sub_cls_name}Warp')
+            )
+          else:
+            # Flattened nested sub-dataclass (Data.contact, Data.efc)
+            _check_fields(cls_name, val, warp_cls, prefix=f'{prefix}{f.name}__')
+          continue
+
+        field_name = f'{prefix}{f.name}'
+        is_batched = getattr(val, '_is_batched', False)
+
+        if is_batched:
+          self.assertTrue(
+              types._BATCH_DIM[cls_name].get(field_name, False),
+              f'Expected {cls_name}.{field_name} with _is_batched=True to be'
+              ' True in _BATCH_DIM',
+          )
+          if cls_name == 'Data':
+            self.assertNotIn(
+                field_name,
+                types.DATA_NON_VMAP,
+                f'Expected batched Data field {field_name} to not be in'
+                ' DATA_NON_VMAP',
+            )
+          ann = warp_cls.__annotations__.get(field_name)
+          if ann is not None:
+            self.assertIn(
+                str(ann),
+                ('jax.Array', "<class 'jax.Array'>"),
+                f'Expected {cls_name}.{field_name} to have jax.Array'
+                ' annotation',
+            )
+        elif isinstance(val, wp.array):
+          self.assertFalse(
+              types._BATCH_DIM[cls_name].get(field_name, True),
+              f'Expected {cls_name}.{field_name} with _is_batched=False to be'
+              ' False in _BATCH_DIM',
+          )
+
+    _check_fields('Model', mw, types.ModelWarp)
+    _check_fields('Data', dw, types.DataWarp)
 
 
 class TileSetTest(parameterized.TestCase):