Merge branch 'google-deepmind:main' into main

doc/changelog.rst

@@ -14,12 +14,30 @@ General
 
 MJX
 ^^^
-3. Added cylinder plane collisions.
+.. admonition:: Breaking API changes
+   :class: attention
+
+   3. Removed deprecated ``mjx.device_get_into`` and ``mjx.device_put`` functions as they lack critical new
+      functionality.
+
+      **Migration:** Use ``mjx.get_data_into`` instead of ``mjx.device_get_into``, and ``mjx.put_data`` instead of
+      ``mjx.device_put``.
+
+4. Added cylinder plane collisions.
+5. Added ``efc_type`` to ``mjx.Data`` and ``dim``, ``efc_address`` to ``mjx.Contact``.
+6. Added ``geom`` to ``mjx.Contact`` and marked ``geom1``, ``geom2`` deprecated.
+7. Added ``ne``, ``nf``, ``nl``, ``nefc``, and ``ncon`` to ``mjx.Data`` to match ``mujoco.MjData``.
+8. Given the above added fields, removed ``mjx.get_params``, ``mjx.ncon``, and ``mjx.count_constraints``.
+9. Changed the way meshes are organized on device to speed up collision detection when a mesh is replicated for many
+   geoms.
+10. Fixed a bug where capsules might be ignored in broadphase colliision checking.
 
 Bug fixes
 ^^^^^^^^^
-4. Defaults of lights were not being saved, now fixed.
+11. Defaults of lights were not being saved, now fixed.
-5. Prevent overwriting of frame names by body names when saving an XML.  Bug introduced in 3.1.4.
+12. Prevent overwriting of frame names by body names when saving an XML.  Introduced in 3.1.4.
+13. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
+    not optional.
 
 
 Version 3.1.4 (April 10th, 2024)

doc/models.rst

@@ -129,3 +129,17 @@ Drones
      - Preview
    * - `Skydio X2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/skydio_x2>`_
      - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/skydio_x2/x2.png
+   * - `Bitcraze Crazyflie 2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/bitcraze_crazyflie_2>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/bitcraze_crazyflie_2/cf2.png
+
+
+Biomechanical
+^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+
+   * - Model
+     - Preview
+   * - `Fruitfly <https://github.com/google-deepmind/mujoco_menagerie/tree/main/flybody>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/flybody/flybody.png

mjx/mujoco/mjx/__init__.py

@@ -16,12 +16,7 @@
 
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
-from mujoco.mjx._src.collision_driver import get_params
-from mujoco.mjx._src.collision_driver import ncon
-from mujoco.mjx._src.constraint import count_constraints
 from mujoco.mjx._src.constraint import make_constraint
-from mujoco.mjx._src.device import device_get_into
-from mujoco.mjx._src.device import device_put
 from mujoco.mjx._src.forward import euler
 from mujoco.mjx._src.forward import forward
 from mujoco.mjx._src.forward import fwd_acceleration

mjx/mujoco/mjx/_src/collision_base.py

@@ -1,67 +0,0 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# 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.
-# ==============================================================================
-"""Collision base."""
-
-import dataclasses
-from typing import Dict, List, Optional, Tuple
-
-import jax
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.dataclasses import PyTreeNode
-from mujoco.mjx._src.types import GeomType
-# pylint: enable=g-importing-member
-
-Contact = Tuple[jax.Array, jax.Array, jax.Array]
-
-
-@dataclasses.dataclass
-class Candidate:
-  geom1: int
-  geom2: int
-  ipair: int
-  geomp: int  # priority geom
-  dim: int
-
-
-CandidateSet = Dict[
-    Tuple[GeomType, GeomType, Tuple[int, ...], Tuple[int, ...]],
-    List[Candidate],
-]
-
-
-class GeomInfo(PyTreeNode):
-  """Collision info for a geom."""
-
-  geom_id: jax.Array
-  pos: jax.Array
-  mat: jax.Array
-  size: jax.Array
-  face: Optional[jax.Array] = None
-  vert: Optional[jax.Array] = None
-  edge_dir: Optional[jax.Array] = None
-  facenorm: Optional[jax.Array] = None
-  edge: Optional[jax.Array] = None
-  edge_face_normal: Optional[jax.Array] = None
-
-
-class SolverParams(PyTreeNode):
-  """Contact solver params."""
-
-  friction: jax.Array
-  solref: jax.Array
-  solreffriction: jax.Array
-  solimp: jax.Array
-  margin: jax.Array
-  gap: jax.Array

mjx/mujoco/mjx/_src/collision_convex.py

@@ -20,12 +20,53 @@ from typing import Tuple
 import jax
 from jax import numpy as jp
 from mujoco.mjx._src import math
+from mujoco.mjx._src import mesh
 # pylint: disable=g-importing-member
-from mujoco.mjx._src.collision_base import Contact
+from mujoco.mjx._src.collision_types import Collision
-from mujoco.mjx._src.collision_base import GeomInfo
+from mujoco.mjx._src.collision_types import ConvexInfo
+from mujoco.mjx._src.collision_types import FunctionKey
+from mujoco.mjx._src.collision_types import GeomInfo
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import GeomType
+from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 
 
+def collider(ncon: int):
+  """Wraps collision functions for use by collision_driver."""
+
+  def wrapper(func):
+    def collide(
+        m: Model, d: Data, key: FunctionKey, geom: jax.Array
+    ) -> Collision:
+      g1, g2 = geom.T
+      infos = [
+          GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]),
+          GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]),
+      ]
+      in_axes = [0, 0]
+      for i in [0, 1]:
+        if key.types[i] == GeomType.BOX:
+          infos[i] = mesh.box(infos[i])
+          in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
+              pos=0, mat=0, face=0, vert=0
+          )
+        elif key.types[i] == GeomType.MESH:
+          infos[i] = mesh.convex(m, key.data_ids[i], infos[i])
+          in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
+              pos=0, mat=0
+          )
+      dist, pos, frame = jax.vmap(func, in_axes=in_axes)(*infos)
+      if ncon > 1:
+        return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
+      return dist, pos, frame
+
+    collide.ncon = ncon
+    return collide
+
+  return wrapper
+
+
 def _closest_segment_point_plane(
     a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array
 ) -> jax.Array:
@@ -178,7 +219,8 @@ def _manifold_points(
   return jp.array([a_idx, b_idx, c_idx, d_idx])
 
 
-def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
+@collider(ncon=4)
+def plane_convex(plane: GeomInfo, convex: ConvexInfo) -> Collision:
   """Calculates contacts between a plane and a convex object."""
   vert = convex.vert
 
@@ -200,10 +242,11 @@ def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
   return dist, pos, frame
 
 
-def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
+@collider(ncon=1)
+def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision:
   """Calculates contact between a sphere and a convex object."""
   faces = convex.face
-  normals = convex.facenorm
+  normals = convex.face_normal
 
   # Put sphere in convex frame.
   sphere_pos = convex.mat.T @ (sphere.pos - convex.pos)
@@ -262,16 +305,15 @@ def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
   n = convex.mat @ n
   pos = convex.mat @ pos + convex.pos
 
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
 
 
-def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
+@collider(ncon=2)
+def capsule_convex(cap: GeomInfo, convex: ConvexInfo) -> Collision:
   """Calculates contacts between a capsule and a convex object."""
   # Get convex transformed normals, faces, and vertices.
   faces = convex.face
-  normals = convex.facenorm
+  normals = convex.face_normal
 
   # Put capsule in convex frame.
   cap_pos = convex.mat.T @ (cap.pos - convex.pos)
@@ -347,7 +389,7 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
       degenerate_edge_dir,
       edge_closest_pt,
       cap_closest_pt,
-  ) = jax.tree_map(lambda x, i=e_idx: jp.take(x, i, axis=0), res)
+  ) = jax.tree_util.tree_map(lambda x, i=e_idx: jp.take(x, i, axis=0), res)
 
   edge_face_normals = edge_face_normal[e_idx]
   edge_voronoi_front = ((edge_face_normals @ edge_axis) < 0).all()
@@ -869,10 +911,9 @@ def _sat_gaussmap(
   return dist, pos, normal
 
 
-def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
+@collider(ncon=4)
+def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision:
   """Calculates contacts between two convex objects."""
-  if c1.face is None or c2.face is None or c1.vert is None or c2.vert is None:
-    raise AssertionError('Mesh info missing.')
   # pad face vertices so that we can broadcast between geom1 and geom2
   # face has shape (n_face, n_vert, 3)
   nvert1, nvert2 = c1.face.shape[1], c2.face.shape[1]
@@ -895,14 +936,14 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
   to_local_mat = c2.mat.T @ c1.mat
 
   faces1 = to_local_pos + faces1 @ to_local_mat.T
-  normals1 = c1.facenorm @ to_local_mat.T
+  normals1 = c1.face_normal @ to_local_mat.T
-  normals2 = c2.facenorm
+  normals2 = c2.face_normal
 
   vertices1 = to_local_pos + c1.vert @ to_local_mat.T
   vertices2 = c2.vert
 
-  unique_edges1 = jp.take(vertices1, c1.edge, axis=0)
+  unique_edges1 = jp.take(vertices1, c1.edge_dir, axis=0)
-  unique_edges2 = jp.take(vertices2, c2.edge, axis=0)
+  unique_edges2 = jp.take(vertices2, c2.edge_dir, axis=0)
 
   edges1 = jp.take(vertices1, c1.edge, axis=0)
   edges2 = jp.take(vertices2, c2.edge, axis=0)
@@ -944,13 +985,6 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
   pos = c2.pos + pos @ c2.mat.T
   normal = normal @ c2.mat.T
   normal = -normal if swapped else normal
-
   frame = jax.vmap(math.make_frame)(normal)
+
   return dist, pos, frame
-
-
-# store ncon as function attributes
-plane_convex.ncon = 4
-sphere_convex.ncon = 1
-capsule_convex.ncon = 2
-convex_convex.ncon = 4

mjx/mujoco/mjx/_src/collision_driver.py

@@ -12,21 +12,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Collide geometries."""
+"""Runs collision checking for all geoms in a Model.
 
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
+To do this, collision_driver builds a collision function table, and then runs
+the collision functions serially on the parameters in the table.
+
+For example, if a Model has three geoms:
+
+geom   |   type
+---------------
+1      | sphere
+2      | capsule
+3      | sphere
+
+collision_driver organizes it into these functions and runs them:
+
+function       | geom pair
+--------------------------
+sphere_sphere  | (1, 3)
+sphere_capsule | (1, 2), (2, 3)
+
+
+Besides collision function, function tables are keyed on mesh id and condim,
+in order to guarantee static shapes for contacts and jacobians.
+"""
+
+import itertools
+from typing import Dict, Iterator, List, Tuple, Union
 
 import jax
 from jax import numpy as jp
 import mujoco
-from mujoco.mjx._src import collision_base
-from mujoco.mjx._src import mesh
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
-from mujoco.mjx._src.collision_base import Candidate
-from mujoco.mjx._src.collision_base import CandidateSet
-from mujoco.mjx._src.collision_base import GeomInfo
-from mujoco.mjx._src.collision_base import SolverParams
 from mujoco.mjx._src.collision_convex import capsule_convex
 from mujoco.mjx._src.collision_convex import convex_convex
 from mujoco.mjx._src.collision_convex import plane_convex
@@ -40,12 +58,14 @@ from mujoco.mjx._src.collision_primitive import sphere_capsule
 from mujoco.mjx._src.collision_primitive import sphere_sphere
 from mujoco.mjx._src.collision_sdf import capsule_ellipsoid
 from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid
+from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
+import numpy as np
 
 # pair-wise collision functions
 _COLLISION_FUNC = {
@@ -70,371 +90,293 @@ _COLLISION_FUNC = {
 }
 
 
-def get_collision_fn(
+# geoms for which we ignore broadphase
-    key: Tuple[Union[GeomType, mujoco.mjtGeom], Union[GeomType, mujoco.mjtGeom]]
+_GEOM_NO_BROADPHASE = {GeomType.HFIELD, GeomType.PLANE}
-) -> Optional[Callable[[GeomInfo, GeomInfo], collision_base.Contact]]:
-  """Returns a collision function given a pair of geom types."""
-  return _COLLISION_FUNC.get(key, None)
 
 
-def _add_candidate(
+def has_collision_fn(t1: GeomType, t2: GeomType) -> bool:
-    result: CandidateSet,
+  """Returns True if a collision function exists for a pair of geom types."""
+  return (t1, t2) in _COLLISION_FUNC
+
+
+def geom_pairs(
     m: Union[Model, mujoco.MjModel],
-    g1: int,
+) -> Iterator[Tuple[int, int, int]]:
-    g2: int,
+  """Returns geom pairs to check for collisions.
-    ipair: int = -1,
-):
-  """Adds a candidate to test for collision."""
-  t1, t2 = m.geom_type[g1], m.geom_type[g2]
-  if t1 > t2:
-    t1, t2, g1, g2 = t2, t1, g2, g1
 
-  # MuJoCo does not collide planes with other planes or hfields
+  Args:
-  if t1 == GeomType.PLANE and t2 == GeomType.PLANE:
+    m: a MuJoCo or MJX model
-    return
-  if t1 == GeomType.PLANE and t2 == GeomType.HFIELD:
-    return
 
-  def mesh_key(i):
+  Yields:
-    convex_data = [[None] * m.ngeom] * 3
+    geom1, geom2, and pair index if defined in <pair> (else -1)
-    if isinstance(m, Model):
+  """
-      convex_data = [
+  pairs = set()
-          m.geom_convex_face,
-          m.geom_convex_vert,
-          m.geom_convex_edge_dir,
-      ]
-    elif isinstance(m, mujoco.MjModel):
-      kwargs = mesh.get(m)
-      convex_data = [
-          kwargs['geom_convex_face'],
-          kwargs['geom_convex_vert'],
-          kwargs['geom_convex_edge_dir'],
-      ]
-    key = tuple((-1,) if v[i] is None else v[i].shape for v in convex_data)
-    return key
 
-  k1, k2 = mesh_key(g1), mesh_key(g2)
+  for i in range(m.npair):
+    g1, g2 = m.pair_geom1[i], m.pair_geom2[i]
+    # order pairs by geom_type for correct function mapping
+    if m.geom_type[g1] > m.geom_type[g2]:
+      g1, g2 = g2, g1
+    pairs.add((g1, g2))
+    yield g1, g2, i
 
-  candidates = {(c.geom1, c.geom2) for c in result.get((t1, t2, k1, k2), [])}
-  if (g1, g2) in candidates:
-    return
-
-  if ipair > -1:
-    candidate = Candidate(g1, g2, ipair, -1, m.pair_dim[ipair])
-  elif m.geom_priority[g1] != m.geom_priority[g2]:
-    gp = g1 if m.geom_priority[g1] > m.geom_priority[g2] else g2
-    candidate = Candidate(g1, g2, -1, gp, m.geom_condim[gp])
-  else:
-    dim = max(m.geom_condim[g1], m.geom_condim[g2])
-    candidate = Candidate(g1, g2, -1, -1, dim)
-
-  result.setdefault((t1, t2, k1, k2), []).append(candidate)
-
-
-def _pair_params(
-    m: Model,
-    candidates: Sequence[Candidate],
-) -> SolverParams:
-  """Gets solver params for pair geoms."""
-  ipair = jp.array([c.ipair for c in candidates])
-  friction = jp.clip(m.pair_friction[ipair], a_min=mujoco.mjMINMU)
-  solref = m.pair_solref[ipair]
-  solreffriction = m.pair_solreffriction[ipair]
-  solimp = m.pair_solimp[ipair]
-  margin = m.pair_margin[ipair]
-  gap = m.pair_gap[ipair]
-
-  return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
-
-
-def _priority_params(
-    m: Model,
-    candidates: Sequence[Candidate],
-) -> SolverParams:
-  """Gets solver params from priority geoms."""
-  geomp = jp.array([c.geomp for c in candidates])
-  friction = m.geom_friction[geomp][:, jp.array([0, 0, 1, 2, 2])]
-  solref = m.geom_solref[geomp]
-  solreffriction = jp.zeros(geomp.shape + (mujoco.mjNREF,))
-  solimp = m.geom_solimp[geomp]
-  g = jp.array([(c.geom1, c.geom2) for c in candidates])
-  margin = jp.amax(m.geom_margin[g.T], axis=0)
-  gap = jp.amax(m.geom_gap[g.T], axis=0)
-
-  return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
-
-
-def _dynamic_params(
-    m: Model,
-    candidates: Sequence[Candidate],
-) -> SolverParams:
-  """Gets solver params for dynamic geoms."""
-  g1 = jp.array([c.geom1 for c in candidates])
-  g2 = jp.array([c.geom2 for c in candidates])
-
-  friction = jp.maximum(m.geom_friction[g1], m.geom_friction[g2])
-  # copy friction terms for the full geom pair
-  friction = friction[:, jp.array([0, 0, 1, 2, 2])]
-
-  minval = jp.array(mujoco.mjMINVAL)
-  solmix1, solmix2 = m.geom_solmix[g1], m.geom_solmix[g2]
-  mix = solmix1 / (solmix1 + solmix2)
-  mix = jp.where((solmix1 < minval) & (solmix2 < minval), 0.5, mix)
-  mix = jp.where((solmix1 < minval) & (solmix2 >= minval), 0.0, mix)
-  mix_fn = jax.vmap(lambda a, b, m: m * a + (1 - m) * b)
-
-  solref1, solref2 = m.geom_solref[g1], m.geom_solref[g2]
-  solref = jp.minimum(solref1, solref2)
-  s_mix = mix_fn(solref1, solref2, mix)
-  solref = jp.where((solref1[0] > 0) & (solref2[0] > 0), s_mix, solref)
-  solreffriction = jp.zeros(g1.shape + (mujoco.mjNREF,))
-  solimp = mix_fn(m.geom_solimp[g1], m.geom_solimp[g2], mix)
-  margin = jp.maximum(m.geom_margin[g1], m.geom_margin[g2])
-  gap = jp.maximum(m.geom_gap[g1], m.geom_gap[g2])
-
-  return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
-
-
-def get_params(
-    m: Union[Model, mujoco.MjModel], candidates: Sequence[Candidate]
-) -> Tuple[List[int], List[int], SolverParams]:
-  """Gets solver params for a list of collision candidates."""
-  # group sol params by different candidate types
-  typ_cands = {}
-  for c in candidates:
-    typ = (c.ipair > -1, c.geomp > -1)
-    typ_cands.setdefault(typ, []).append(c)
-
-  geom1, geom2, params = [], [], []
-  for (pair, priority), candidates in typ_cands.items():
-    geom1.extend([c.geom1 for c in candidates])
-    geom2.extend([c.geom2 for c in candidates])
-    if pair:
-      params.append(_pair_params(m, candidates))
-    elif priority:
-      params.append(_priority_params(m, candidates))
-    else:
-      params.append(_dynamic_params(m, candidates))
-
-  params = jax.tree_map(lambda *x: jp.concatenate(x), *params)
-  return geom1, geom2, params
-
-
-def _pair_info(
-    m: Model, d: Data, geom1: Sequence[int], geom2: Sequence[int]
-) -> Tuple[GeomInfo, GeomInfo, Sequence[Dict[str, Optional[int]]]]:
-  """Returns geom pair info for calculating collision."""
-  def mesh_info(geom):
-    g = jp.array(geom)
-    info = GeomInfo(
-        g,
-        d.geom_xpos[g],
-        d.geom_xmat[g],
-        m.geom_size[g],
-    )
-    in_axes = jax.tree_map(lambda x: 0, info)
-    is_mesh = m.geom_convex_face[geom[0]] is not None
-    if is_mesh:
-      info = info.replace(
-          face=jp.stack([m.geom_convex_face[i] for i in geom]),
-          vert=jp.stack([m.geom_convex_vert[i] for i in geom]),
-          edge_dir=jp.stack([m.geom_convex_edge_dir[i] for i in geom]),
-          facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom]),
-          edge=jp.stack([m.geom_convex_edge[i] for i in geom]),
-          edge_face_normal=jp.stack(
-              [m.geom_convex_edge_face_normal[i] for i in geom]
-          ),
-      )
-      in_axes = in_axes.replace(
-          face=0,
-          vert=0,
-          edge_dir=0,
-          facenorm=0,
-          edge=0,
-          edge_face_normal=0,
-      )
-    return info, in_axes
-
-  info1, in_axes1 = mesh_info(geom1)
-  info2, in_axes2 = mesh_info(geom2)
-  return info1, info2, [in_axes1, in_axes2]
-
-
-def _body_pair_filter(
-    m: Union[Model, mujoco.MjModel], b1: int, b2: int
-) -> bool:
-  """Filters body pairs for collision."""
-  dsbl_filterparent = m.opt.disableflags & DisableBit.FILTERPARENT
-  weld1 = m.body_weldid[b1]
-  weld2 = m.body_weldid[b2]
-  parent_weld1 = m.body_weldid[m.body_parentid[weld1]]
-  parent_weld2 = m.body_weldid[m.body_parentid[weld2]]
-
-  if weld1 == weld2:
-    # filter out self-collisions
-    return True
-
-  if (
-      not dsbl_filterparent
-      and weld1 != 0
-      and weld2 != 0
-      and (weld1 == parent_weld2 or weld2 == parent_weld1)
-  ):
-    # filter out parent-child collisions
-    return True
-
-  return False
-
-
-def _broadphase_enabled(
-    geom_types: Tuple[GeomType, GeomType],
-    n_pairs: int,
-    max_pairs: int,
-) -> bool:
-  return (
-      GeomType.PLANE not in geom_types
-      and max_pairs > -1
-      and n_pairs > max_pairs
-  )
-
-
-def _collide_geoms(
-    m: Model,
-    d: Data,
-    geom_types: Tuple[GeomType, GeomType],
-    candidates: Sequence[Candidate],
-) -> Contact:
-  """Collides a geom pair."""
-  fn = get_collision_fn(geom_types)
-  if not fn:
-    return Contact.zero()
-
-  geom1, geom2, params = get_params(m, candidates)
-  g1, g2, in_axes = _pair_info(m, d, geom1, geom2)
-
-  # Run a crude version of broadphase.
-  max_pairs = int(support.get_custom_numeric(m, 'max_geom_pairs'))
-  run_broadphase = _broadphase_enabled(geom_types, len(geom1), max_pairs)
-  n_pairs = max_pairs if run_broadphase else len(geom1)
-  if run_broadphase:
-    # broadphase over geom pairs, using bounding spheres
-    size1 = jp.max(m.geom_size[g1.geom_id], axis=-1)
-    size2 = jp.max(m.geom_size[g2.geom_id], axis=-1)
-    dists = jax.vmap(jp.linalg.norm)(g2.pos - g1.pos) - (size1 + size2)
-    _, idx = jax.lax.top_k(-dists, k=n_pairs)
-    g1, g2, params = jax.tree_map(
-        lambda x, idx=idx: x[idx, ...], (g1, g2, params)
-    )
-
-  # call contact function
-  res = jax.vmap(fn, in_axes=in_axes)(g1, g2)
-  dist, pos, frame = jax.tree_map(jp.concatenate, res)
-
-  # repeat params by the number of contacts per geom pair
-  geom1, geom2, params = jax.tree_map(
-      lambda x: jp.repeat(x, fn.ncon, axis=0),  # pytype: disable=attribute-error
-      (g1.geom_id, g2.geom_id, params),
-  )
-
-  con = Contact(
-      dist=dist,
-      pos=pos,
-      frame=frame,
-      includemargin=params.margin - params.gap,
-      friction=params.friction,
-      solref=params.solref,
-      solreffriction=params.solreffriction,
-      solimp=params.solimp,
-      geom1=geom1,
-      geom2=geom2,
-  )
-  return con
-
-
-def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
-  """Returns candidates for collision checking."""
-  candidate_set = {}
-
-  for ipair in range(m.npair):
-    g1, g2 = m.pair_geom1[ipair], m.pair_geom2[ipair]
-    _add_candidate(candidate_set, m, g1, g2, ipair)
-
-  body_pairs = []
   exclude_signature = set(m.exclude_signature)
   geom_con = m.geom_contype | m.geom_conaffinity
+  filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT)
   b_start = m.body_geomadr
   b_end = b_start + m.body_geomnum
 
   for b1 in range(m.nbody):
     if not geom_con[b_start[b1]:b_end[b1]].any():
       continue
+    w1 = m.body_weldid[b1]
+    w1_p = m.body_weldid[m.body_parentid[w1]]
+
     for b2 in range(b1, m.nbody):
       if not geom_con[b_start[b2]:b_end[b2]].any():
         continue
       signature = (b1 << 16) + (b2)
       if signature in exclude_signature:
         continue
-      if _body_pair_filter(m, b1, b2):
+      w2 = m.body_weldid[b2]
+      # ignore self-collisions
+      if w1 == w2:
         continue
-      body_pairs.append((b1, b2))
+      w2_p = m.body_weldid[m.body_parentid[w2]]
+      # ignore parent-child collisions
+      if filterparent and w1 != 0 and w2 != 0 and (w1 == w2_p or w2 == w1_p):
+        continue
+      g1_range = [g for g in range(b_start[b1], b_end[b1]) if geom_con[g]]
+      g2_range = [g for g in range(b_start[b2], b_end[b2]) if geom_con[g]]
 
-  for b1, b2 in body_pairs:
+      for g1, g2 in itertools.product(g1_range, g2_range):
-    for g1 in range(b_start[b1], b_end[b1]):
+        t1, t2 = m.geom_type[g1], m.geom_type[g2]
-      if not geom_con[g1]:
+        # order pairs by geom_type for correct function mapping
-        continue
+        if t1 > t2:
-      for g2 in range(b_start[b2], b_end[b2]):
+          g1, g2, t1, t2 = g2, g1, t2, t1
-        if not geom_con[g2]:
+        # ignore plane<>plane and plane<>hfield
+        if (t1, t2) == (GeomType.PLANE, GeomType.PLANE):
           continue
+        if (t1, t2) == (GeomType.PLANE, GeomType.HFIELD):
+          continue
+        # geoms must match contype and conaffinity on some bit
         mask = m.geom_contype[g1] & m.geom_conaffinity[g2]
         mask |= m.geom_contype[g2] & m.geom_conaffinity[g1]
-        if mask != 0:
+        if not mask:
-          _add_candidate(candidate_set, m, g1, g2)
+          continue
 
-  return candidate_set
+        if (g1, g2) not in pairs:
+          pairs.add((g1, g2))
+          yield g1, g2, -1
 
 
-def ncon(m: Union[Model, mujoco.MjModel]) -> int:
+def _geom_groups(
-  """Returns the number of contacts computed in MJX given a model."""
+    m: Union[Model, mujoco.MjModel],
+) -> Dict[FunctionKey, List[Tuple[int, int, int]]]:
+  """Returns geom pairs to check for collision grouped by collision function.
+
+  The grouping consists of:
+    - The collision function to run, which is determined by geom types
+    - For mesh geoms, convex functions are run for each distinct mesh in the
+      model, because the convex functions expect static mesh size. If a sphere
+      collides with a cube and a tetrahedron, sphere_convex is called twice.
+    - The condim of the collision. This ensures that the size of the resulting
+      constraint jacobian is determined at compile time.
+
+  Args:
+    m: a MuJoCo or MJX model
+
+  Returns:
+    a dict with grouping key and values geom1, geom2, pair index
+  """
+  groups = {}
+
+  for g1, g2, ip in geom_pairs(m):
+    types = m.geom_type[g1], m.geom_type[g2]
+    data_ids = m.geom_dataid[g1], m.geom_dataid[g2]
+    if ip > -1:
+      condim = m.pair_dim[ip]
+    elif m.geom_priority[g1] > m.geom_priority[g2]:
+      condim = m.geom_condim[g1]
+    elif m.geom_priority[g1] < m.geom_priority[g2]:
+      condim = m.geom_condim[g2]
+    else:
+      condim = max(m.geom_condim[g1], m.geom_condim[g2])
+
+    key = FunctionKey(types, data_ids, condim)
+    groups.setdefault(key, []).append((g1, g2, ip))
+
+  return groups
+
+
+def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
+  """Returns contact groups to check for collisions.
+
+  Contacts are grouped the same way as _geom_groups.  Only one contact is
+  emitted per geom pair, even if the collision function emits multiple contacts.
+
+  Args:
+    m: MJX model
+    d: MJX data
+
+  Returns:
+    a dict where the key is the grouping and value is a Contact
+  """
+  groups = {}
+  eps = mujoco.mjMINVAL
+
+  for key, geom_ids in _geom_groups(m).items():
+    geom = np.array(geom_ids)
+    geom1, geom2, ip = geom.T
+    geom1, geom2, ip = geom1[ip == -1], geom2[ip == -1], ip[ip != -1]
+    params = []
+
+    if ip.size > 0:
+      # pair contacts get their params from m.pair_* fields
+      params.append((
+          m.pair_margin[ip] - m.pair_gap[ip],
+          jp.clip(m.pair_friction[ip], a_min=eps),
+          m.pair_solref[ip],
+          m.pair_solreffriction[ip],
+          m.pair_solimp[ip]
+      ))
+    if geom1.size > 0 and geom2.size > 0:
+      # other contacts get their params from geom fields
+      margin = jp.maximum(m.geom_margin[geom1], m.geom_margin[geom2])
+      gap = jp.maximum(m.geom_gap[geom1], m.geom_gap[geom2])
+      solmix1, solmix2 = m.geom_solmix[geom1], m.geom_solmix[geom2]
+      mix = solmix1 / (solmix1 + solmix2)
+      mix = jp.where((solmix1 < eps) & (solmix2 < eps), 0.5, mix)
+      mix = jp.where((solmix1 < eps) & (solmix2 >= eps), 0.0, mix)
+      mix = jp.where((solmix1 >= eps) & (solmix2 < eps), 1.0, mix)
+      mix = mix[:, None]  # for correct broadcasting
+      # friction: max
+      friction = jp.maximum(m.geom_friction[geom1], m.geom_friction[geom2])
+      solref1, solref2 = m.geom_solref[geom1], m.geom_solref[geom2]
+      # reference standard: mix
+      solref_standard = mix * solref1 + (1 - mix) * solref2
+      # reference direct: min
+      solref_direct = jp.minimum(solref1, solref2)
+      is_standard = (solref1[:, [0, 0]] > 0) & (solref2[:, [0, 0]] > 0)
+      solref = jp.where(is_standard, solref_standard, solref_direct)
+      solreffriction = jp.zeros(geom1.shape + (mujoco.mjNREF,))
+      # impedance: mix
+      solimp = mix * m.geom_solimp[geom1] + (1 - mix) * m.geom_solimp[geom2]
+
+      pri = m.geom_priority[geom1] != m.geom_priority[geom2]
+      if pri.any():
+        # use priority geom when specified instead of mixing
+        gp1, gp2 = m.geom_priority[geom1], m.geom_priority[geom2]
+        gp = np.where(gp1 > gp2, geom1, geom2)[pri]
+        friction = friction.at[pri].set(m.geom_friction[gp])
+        solref = solref.at[pri].set(m.geom_solref[gp])
+        solimp = solimp.at[pri].set(m.geom_solimp[gp])
+
+      # unpack 5d friction:
+      friction = friction[:, [0, 0, 1, 2, 2]]
+      params.append((margin - gap, friction, solref, solreffriction, solimp))
+
+    params = map(jp.concatenate, zip(*params))
+    includemargin, friction, solref, solreffriction, solimp = params
+
+    groups[key] = Contact(
+        # dist, pos, frame get filled in by collision functions:
+        dist=None,
+        pos=None,
+        frame=None,
+        includemargin=includemargin,
+        friction=friction,
+        solref=solref,
+        solreffriction=solreffriction,
+        solimp=solimp,
+        dim=d.contact.dim,
+        geom1=jp.array(geom[:, 0]),
+        geom2=jp.array(geom[:, 1]),
+        geom=jp.array(geom[:, :2]),
+        efc_address=d.contact.efc_address,
+    )
+
+  return groups
+
+
+def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
+  """Returns the dims of the contacts for a Model."""
   if m.opt.disableflags & DisableBit.CONTACT:
-    return 0
+    return np.empty(0, dtype=int)
 
-  candidates = collision_candidates(m)
+  group_counts = {k: len(v) for k, v in _geom_groups(m).items()}
-  max_count = int(support.get_custom_numeric(m, 'max_contact_points'))
-  max_pairs = int(support.get_custom_numeric(m, 'max_geom_pairs'))
 
-  count = 0
+  # max_geom_pairs limits the number of pairs we process in a collision function
-  for k, v in candidates.items():
+  # by first running a primitive broad phase culling on the pairs
-    fn = get_collision_fn(k[0:2])
+  max_geom_pairs = support.get_custom_int(m, 'max_geom_pairs')
-    if fn is None:
+  if max_geom_pairs > -1:
-      continue
+    for k in group_counts:
-    run_broadphase = _broadphase_enabled((k[0], k[1]), len(v), max_pairs)
+      if set(k.types) & _GEOM_NO_BROADPHASE:
-    n_pair = max_pairs if run_broadphase else len(v)
+        continue
-    count += n_pair * fn.ncon  # pytype: disable=attribute-error
+      group_counts[k] = min(group_counts[k], max_geom_pairs)
 
-  return min(max_count, count) if max_count > -1 else count
+  # max_contact_points limits the number of contacts emitted by selecting the
+  # contacts with the most penetration after calling collision functions
+  max_contact_points = support.get_custom_int(m, 'max_contact_points')
+
+  condim_counts = {}
+  for k, v in group_counts.items():
+    func = _COLLISION_FUNC[k.types]
+    num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v  # pytype: disable=attribute-error
+    if max_contact_points > -1:
+      num_contacts = min(max_contact_points, num_contacts)
+    condim_counts[k.condim] = num_contacts
+
+  dims = sum(([c] * condim_counts[c] for c in sorted(condim_counts)), [])
+
+  return np.array(dims)
 
 
 def collision(m: Model, d: Data) -> Data:
   """Collides geometries."""
-  if ncon(m) == 0:
+  if d.ncon == 0:
-    return d.replace(contact=Contact.zero())
+    return d
 
-  candidate_set = collision_candidates(m)
+  groups = _contact_groups(m, d)
+  max_geom_pairs = support.get_custom_int(m, 'max_geom_pairs')
+  max_contact_points = support.get_custom_int(m, 'max_contact_points')
 
-  contacts = []
+  # run collision functions on groups
-  for key, candidates in candidate_set.items():
+  for key, contact in groups.items():
-    geom_types = key[0:2]
+    # determine which contacts we'll use for collision testing by running a
-    contacts.append(_collide_geoms(m, d, geom_types, candidates))
+    # broad phase cull if requested
+    if max_geom_pairs > -1 and contact.geom.shape[0] > max_geom_pairs:
+      pos1, pos2 = d.geom_xpos[contact.geom.T]
+      size1, size2 = m.geom_rbound[contact.geom.T]
+      dist = jax.vmap(jp.linalg.norm)(pos2 - pos1) - (size1 + size2)
+      _, idx = jax.lax.top_k(-dist, k=max_geom_pairs)
+      contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact)
 
-  if not contacts:
+    # run the collision function specified by the grouping key
-    raise RuntimeError('No contacts found.')
+    func = _COLLISION_FUNC[key.types]
+    dist, pos, frame = func(m, d, key, contact.geom)
+    ncon = func.ncon  # pytype: disable=attribute-error
+    if ncon > 1:
+      # repeat contacts to match the number of collisions returned
+      repeat_fn = lambda x, r=ncon: jp.repeat(x, r, axis=0)
+      contact = jax.tree_util.tree_map(repeat_fn, contact)
+    groups[key] = contact.replace(dist=dist, pos=pos, frame=frame)
 
-  contact = jax.tree_map(lambda *x: jp.concatenate(x), *contacts)
+  # collapse contacts together, ensuring they are grouped by condim
+  condim_groups = {}
+  for key, contact in groups.items():
+    condim_groups.setdefault(key.condim, []).append(contact)
 
-  max_contact_points = int(support.get_custom_numeric(m, 'max_contact_points'))
+  # limit the number of contacts per condim group if requested
-  if max_contact_points > -1 and contact.dist.shape[0] > max_contact_points:
+  if max_contact_points > -1:
-    # get top-k contacts
+    for key, contacts in condim_groups.items():
-    _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points)
+      contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)
-    contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact)
+      if contact.geom.shape[0] > max_contact_points:
+        _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points)
+        contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact)
+      condim_groups[key] = [contact]
+
+  contacts = sum([condim_groups[k] for k in sorted(condim_groups)], [])
+  contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)
 
   return d.replace(contact=contact)

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -402,7 +402,8 @@ class CapsuleCollisionTest(parameterized.TestCase):
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
 
@@ -437,7 +438,8 @@ class CapsuleCollisionTest(parameterized.TestCase):
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip1', 1e-4)
     np.testing.assert_array_almost_equal(
@@ -457,7 +459,8 @@ class CapsuleCollisionTest(parameterized.TestCase):
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip2', 1e-4)
     np.testing.assert_array_almost_equal(
@@ -494,7 +497,8 @@ class CylinderTest(absltest.TestCase):
     d.contact.pos[:] = d.contact.pos[idx]
 
     # extract the contact points with penetration
-    c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:2], dx.contact)
+    c = c.replace(dim=c.dim[:2], efc_address=c.efc_address[:2])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'cylinder_plane', 1e-5)
 
@@ -531,7 +535,8 @@ class ConvexTest(absltest.TestCase):
     np.testing.assert_array_less(dx.contact.dist[:2], 0)
     np.testing.assert_array_less(-dx.contact.dist[2:], 0)
     # extract the contact points with penetration
-    c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
+    c = jax.tree_util.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
+    c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-5)
 
@@ -615,7 +620,8 @@ class ConvexTest(absltest.TestCase):
     np.testing.assert_array_less(dx.contact.dist[:1], 0)
     np.testing.assert_array_less(-dx.contact.dist[1:], 0)
     # extract the contact point with penetration
-    c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
 
@@ -763,28 +769,28 @@ class BodyPairFilterTest(absltest.TestCase):
     self.assertEqual(dx.contact.pos.shape[0], 1)
 
 
-class NconTest(parameterized.TestCase):
+class DimTest(parameterized.TestCase):
-  """Tests ncon."""
+  """Tests contact dim."""
 
   def test_ncon(self):
     m = test_util.load_test_file('constraints.xml')
-    ncon = collision_driver.ncon(m)
+    dim = collision_driver.make_condim(m)
-    self.assertEqual(ncon, 16)
+    np.testing.assert_array_equal(dim, np.array([3] * 16))
 
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags |= DisableBit.CONTACT
-    ncon = collision_driver.ncon(m)
+    dim = collision_driver.make_condim(m)
-    self.assertEqual(ncon, 0)
+    self.assertEqual(dim.size, 0)
 
   def test_ncon_meshes(self):
     m = test_util.load_test_file('shadow_hand/scene_right.xml')
 
-    ncon = collision_driver.ncon(m)
+    ncon = collision_driver.make_condim(m).size
     self.assertEqual(ncon, 15)
 
     mx = mjx.put_model(m)
-    ncon = collision_driver.ncon(mx)
+    ncon = collision_driver.make_condim(mx).size
     self.assertEqual(ncon, 15)
 
     # get rid of max_contact_points, test only max_geom_pairs
@@ -795,11 +801,11 @@ class NconTest(parameterized.TestCase):
       if name_ == 'max_contact_points':
         m.numeric_data[m.numeric_adr[i]] = -1
 
-    ncon = collision_driver.ncon(m)
+    ncon = collision_driver.make_condim(m).size
     self.assertEqual(ncon, 98)
 
     mx = mjx.put_model(m)
-    ncon = collision_driver.ncon(mx)
+    ncon = collision_driver.make_condim(mx).size
     self.assertEqual(ncon, 98)
 
 

mjx/mujoco/mjx/_src/collision_primitive.py

@@ -20,34 +20,53 @@ import jax
 from jax import numpy as jp
 from mujoco.mjx._src import math
 # pylint: disable=g-importing-member
-from mujoco.mjx._src.collision_base import Contact
+from mujoco.mjx._src.collision_types import Collision
-from mujoco.mjx._src.collision_base import GeomInfo
+from mujoco.mjx._src.collision_types import GeomInfo
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 
 
+def collider(ncon: int):
+  """Wraps collision functions for use by collision_driver."""
+  def wrapper(func):
+    def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
+      g1, g2 = geom.T
+      info1 = GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1])
+      info2 = GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2])
+      dist, pos, frame = jax.vmap(func)(info1, info2)
+      if ncon > 1:
+        return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
+      return dist, pos, frame
+
+    collide.ncon = ncon
+    return collide
+
+  return wrapper
+
+
 def _plane_sphere(
     plane_normal: jax.Array,
     plane_pos: jax.Array,
     sphere_pos: jax.Array,
-    radius: jax.Array,
+    sphere_radius: jax.Array,
 ) -> Tuple[jax.Array, jax.Array]:
-  """Returns the penetration and contact point between a plane and sphere."""
+  """Returns the distance and contact point between a plane and sphere."""
-  cdist = jp.dot(sphere_pos - plane_pos, plane_normal)
+  dist = jp.dot(sphere_pos - plane_pos, plane_normal) - sphere_radius
-  dist = cdist - radius
+  pos = sphere_pos - plane_normal * (sphere_radius + 0.5 * dist)
-  pos = sphere_pos - plane_normal * (radius + 0.5 * dist)
   return dist, pos
 
 
-def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Contact:
+@collider(ncon=1)
+def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Collision:
   """Calculates contact between a plane and a sphere."""
   n = plane.mat[:, 2]
   dist, pos = _plane_sphere(n, plane.pos, sphere.pos, sphere.size[0])
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
 
 
-def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact:
+@collider(ncon=2)
+def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Collision:
   """Calculates two contacts between a capsule and a plane."""
   n, axis = plane.mat[:, 2], cap.mat[:, 2]
   # align contact frames with capsule axis
@@ -56,16 +75,17 @@ def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact:
   b = jp.where(b_norm < 0.5, jp.where((-0.5 < n[1]) & (n[1] < 0.5), y, z), b)
   frame = jp.array([[n, b, jp.cross(n, b)]])
   segment = axis * cap.size[1]
-  contacts = []
+  collisions = []
   for offset in [segment, -segment]:
     dist, pos = _plane_sphere(n, plane.pos, cap.pos + offset, cap.size[0])
     dist = jp.expand_dims(dist, axis=0)
     pos = jp.expand_dims(pos, axis=0)
-    contacts.append((dist, pos, frame))
+    collisions.append((dist, pos, frame))
-  return jax.tree_map(lambda *x: jp.concatenate(x), *contacts)
+  return jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *collisions)
 
 
-def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Contact:
+@collider(ncon=1)
+def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Collision:
   """Calculates one contact between an ellipsoid and a plane."""
   n = plane.mat[:, 2]
   size = ellipsoid.size
@@ -73,12 +93,11 @@ def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Contact:
   pos = ellipsoid.pos + ellipsoid.mat @ (sphere_support * size)
   dist = jp.dot(n, pos - plane.pos)
   pos = pos - n * dist * 0.5
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
 
 
-def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Contact:
+@collider(ncon=3)
+def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision:
   """Calculates one contact between an cylinder and a plane."""
   n = plane.mat[:, 2]
   axis = cylinder.mat[:, 2]
@@ -139,7 +158,7 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Contact:
 
 def _sphere_sphere(
     pos1: jax.Array, radius1: jax.Array, pos2: jax.Array, radius2: jax.Array
-) -> Contact:
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
   """Returns the penetration, contact point, and normal between two spheres."""
   n, dist = math.normalize_with_norm(pos2 - pos1)
   n = jp.where(dist == 0.0, jp.array([1.0, 0.0, 0.0]), n)
@@ -148,15 +167,15 @@ def _sphere_sphere(
   return dist, pos, n
 
 
-def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Contact:
+@collider(ncon=1)
+def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Collision:
   """Calculates contact between two spheres."""
   dist, pos, n = _sphere_sphere(s1.pos, s1.size[0], s2.pos, s2.size[0])
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
 
 
-def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact:
+@collider(ncon=1)
+def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Collision:
   """Calculates one contact between a sphere and a capsule."""
   axis, length = cap.mat[:, 2], cap.size[1]
   segment = axis * length
@@ -164,19 +183,14 @@ def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact:
       cap.pos - segment, cap.pos + segment, sphere.pos
   )
   dist, pos, n = _sphere_sphere(sphere.pos, sphere.size[0], pt, cap.size[0])
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
 
 
-def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact:
+@collider(ncon=1)
+def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Collision:
   """Calculates one contact between two capsules."""
-  axis1, length1, axis2, length2 = (
+  axis1, length1 = cap1.mat[:, 2], cap1.size[1]
-      cap1.mat[:, 2],
+  axis2, length2 = cap2.mat[:, 2], cap2.size[1]
-      cap1.size[1],
-      cap2.mat[:, 2],
-      cap2.size[1],
-  )
   seg1, seg2 = axis1 * length1, axis2 * length2
   pt1, pt2 = math.closest_segment_to_segment_points(
       cap1.pos - seg1,
@@ -186,15 +200,4 @@ def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact:
   )
   radius1, radius2 = cap1.size[0], cap2.size[0]
   dist, pos, n = _sphere_sphere(pt1, radius1, pt2, radius2)
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
-
-# store ncon as function attributes
-plane_sphere.ncon = 1
-plane_capsule.ncon = 2
-plane_ellipsoid.ncon = 1
-plane_cylinder.ncon = 3
-sphere_sphere.ncon = 1
-sphere_capsule.ncon = 1
-capsule_capsule.ncon = 1

mjx/mujoco/mjx/_src/collision_sdf.py

@@ -28,15 +28,35 @@ import jax
 from jax import numpy as jp
 from mujoco.mjx._src import math
 # pylint: disable=g-importing-member
-from mujoco.mjx._src.collision_base import Contact
+from mujoco.mjx._src.collision_types import Collision
-from mujoco.mjx._src.collision_base import GeomInfo
+from mujoco.mjx._src.collision_types import GeomInfo
 from mujoco.mjx._src.dataclasses import PyTreeNode
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 
 # the SDF function takes position in, and returns a distance or objective
 SDFFn = Callable[[jax.Array], jax.Array]
 
 
+def collider(ncon: int):
+  """Wraps collision functions for use by collision_driver."""
+  def wrapper(func):
+    def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
+      g1, g2 = geom.T
+      info1 = GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1])
+      info2 = GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2])
+      dist, pos, frame = jax.vmap(func)(info1, info2)
+      if ncon > 1:
+        return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
+      return dist, pos, frame
+
+    collide.ncon = ncon
+    return collide
+
+  return wrapper
+
+
 def _plane(pos: jax.Array, size: jax.Array) -> jax.Array:
   del size
   return pos[2]
@@ -128,22 +148,15 @@ def _optim(
   return pos, dist, n
 
 
-def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Contact:
+@collider(ncon=1)
+def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision:
   """"Calculates contact between a capsule and an ellipsoid."""
   pos, dist, n = _optim(_capsule, _ellipsoid, c, e)
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
 
 
-def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Contact:
+@collider(ncon=1)
+def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision:
   """"Calculates contact between two ellipsoids."""
   pos, dist, n = _optim(_ellipsoid, _ellipsoid, e1, e2)
-  return jax.tree_map(
+  return dist, pos, math.make_frame(n)
-      lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
-  )
-
-
-# store ncon as function attributes
-capsule_ellipsoid.ncon = 1
-ellipsoid_ellipsoid.ncon = 1

mjx/mujoco/mjx/_src/collision_types.py

@@ -0,0 +1,65 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# 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.
+# ==============================================================================
+"""Collision base types."""
+
+import dataclasses
+from typing import Tuple
+import jax
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.dataclasses import PyTreeNode
+# pylint: enable=g-importing-member
+
+
+class GeomInfo(PyTreeNode):
+  """Geom propertes of primitive and SDF shapes."""
+  pos: jax.Array
+  mat: jax.Array
+  size: jax.Array
+
+
+class ConvexInfo(PyTreeNode):
+  """Geom propertes of convex meshes."""
+  pos: jax.Array
+  mat: jax.Array
+  vert: jax.Array
+  face: jax.Array
+  face_normal: jax.Array
+  edge: jax.Array
+  edge_face_normal: jax.Array
+  edge_dir: jax.Array
+
+
+# Collision returned by collision functions:
+#  - distance          distance between nearest points; neg: penetration
+#  - position  (3,)    position of contact point: midpoint between geoms
+#  - frame     (3, 3)  normal is in [0, :], points from geom[0] to geom[1]
+Collision = Tuple[jax.Array, jax.Array, jax.Array]
+
+
+@dataclasses.dataclass(frozen=True)
+class FunctionKey:
+  """Specifies how geom pairs group into collision_driver's function table.
+
+  Attributes:
+    types: geom type pair, which determines the collision function
+    data_ids: geom data id pair: mesh id for mesh geoms, otherwise -1.
+              Meshes have distinct face/vertex counts, so must occupy distinct
+              entries in the collision function table.
+    condim: grouping by condim of the colliision ensures that the size of the
+            resulting constraint jacobian is determined at compile time.
+  """
+  types: Tuple[int, int]
+  data_ids: Tuple[int, int]
+  condim: int

mjx/mujoco/mjx/_src/constraint.py

@@ -24,6 +24,7 @@ from mujoco.mjx._src import math
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.dataclasses import PyTreeNode
+from mujoco.mjx._src.types import ConstraintType
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
@@ -34,6 +35,9 @@ from mujoco.mjx._src.types import Model
 import numpy as np
 
 
+_CONDIM_EFC_COUNT = {1: 1, 3: 4, 4: 6, 6: 10}
+
+
 class _Efc(PyTreeNode):
   """Support data for creating constraint matrices."""
   J: jax.Array
@@ -111,7 +115,7 @@ def _instantiate_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
     return j, cpos, jp.repeat(math.norm(cpos), 3)
 
   # concatenate to drop connect grouping dimension
-  j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2))
+  j, pos, pos_norm = jax.tree_util.tree_map(jp.concatenate, fn(data, id1, id2))
   invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0]
   invweight = jp.repeat(invweight, 3)
   solref = jp.tile(m.eq_solref[ids], (3, 1))
@@ -164,7 +168,7 @@ def _instantiate_equality_weld(m: Model, d: Data) -> Optional[_Efc]:
     return j, pos, jp.repeat(math.norm(pos), 6)
 
   # concatenate to drop weld grouping dimension
-  j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2))
+  j, pos, pos_norm = jax.tree_util.tree_map(jp.concatenate, fn(data, id1, id2))
   invweight = m.body_invweight0[id1] + m.body_invweight0[id2]
   invweight = jp.repeat(invweight, 3)
   solref = jp.tile(m.eq_solref[ids], (6, 1))
@@ -277,7 +281,7 @@ def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
 def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for contacts."""
 
-  if collision_driver.ncon(m) == 0:
+  if d.ncon == 0:
     return None
 
   @jax.vmap
@@ -308,42 +312,57 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
 
   res = fn(d.contact)
   # remove contact grouping dimension:
-  j, invweight, pos, solref, solimp = jax.tree_map(jp.concatenate, res)
+  j, invweight, pos, solref, solimp = jax.tree_util.tree_map(jp.concatenate, res)
   frictionloss = jp.zeros_like(pos)
 
   return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
 
 
-def count_constraints(
+def counts(efc_type: np.ndarray) -> Tuple[int, int, int, int]:
-    m: Union[Model, mujoco.MjModel], d: Optional[Data] = None
-) -> Tuple[int, int, int, int]:
   """Returns equality, friction, limit, and contact constraint counts."""
-  if m.opt.disableflags & DisableBit.CONSTRAINT:
+  ne = (efc_type == ConstraintType.EQUALITY).sum()
-    return 0, 0, 0, 0
+  nf = 0  # no support for friction loss yet
-
+  nl = (efc_type == ConstraintType.LIMIT_JOINT).sum()
-  if m.opt.disableflags & DisableBit.EQUALITY:
+  nc = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum()
-    ne = 0
-  else:
-    ne_connect = (m.eq_type == EqType.CONNECT).sum()
-    ne_weld = (m.eq_type == EqType.WELD).sum()
-    ne_joint = (m.eq_type == EqType.JOINT).sum()
-    ne = ne_connect * 3 + ne_weld * 6 + ne_joint
-
-  nf = 0
-
-  if m.opt.disableflags & DisableBit.LIMIT:
-    nl = 0
-  else:
-    nl = int(m.jnt_limited.sum())
-
-  if d is None:
-    nc = collision_driver.ncon(m) * 4
-  else:
-    nc = d.efc_J.shape[-2] - ne - nf - nl
 
   return ne, nf, nl, nc
 
 
+def make_efc_type(
+    m: Union[Model, mujoco.MjModel], dim: Optional[np.ndarray] = None
+) -> np.ndarray:
+  """Returns efc_type that outlines the type of each constraint row."""
+  if m.opt.disableflags & DisableBit.CONSTRAINT:
+    return np.empty(0, dtype=int)
+
+  dim = collision_driver.make_condim(m) if dim is None else dim
+  efc_types = []
+
+  if not m.opt.disableflags & DisableBit.EQUALITY:
+    num_rows = (m.eq_type == EqType.CONNECT).sum() * 3
+    num_rows += (m.eq_type == EqType.WELD).sum() * 6
+    num_rows += (m.eq_type == EqType.JOINT).sum()
+    efc_types.extend([ConstraintType.EQUALITY] * num_rows)
+
+  if not m.opt.disableflags & DisableBit.LIMIT:
+    efc_types.extend([ConstraintType.LIMIT_JOINT] * m.jnt_limited.sum())
+
+  if not m.opt.disableflags & DisableBit.CONTACT:
+    num_rows = sum(_CONDIM_EFC_COUNT[d] for d in dim)
+    efc_types.extend([ConstraintType.CONTACT_PYRAMIDAL] * num_rows)
+
+  return np.array(efc_types)
+
+
+def make_efc_address(efc_type: np.ndarray, dim: np.ndarray) -> np.ndarray:
+  """Returns efc_address that maps contacts to constraint row address."""
+  nc = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum()
+  nc_start = efc_type.size - nc
+  offsets = np.cumsum([0] + [_CONDIM_EFC_COUNT[d] for d in dim])[:-1]
+
+  return nc_start + offsets
+
+
 def make_constraint(m: Model, d: Data) -> Data:
   """Creates constraint jacobians and other supporting data."""
 
@@ -366,7 +385,7 @@ def make_constraint(m: Model, d: Data) -> Data:
     d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z)
     return d
 
-  efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
+  efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
 
   @jax.vmap
   def fn(efc):

mjx/mujoco/mjx/_src/constraint_test.py

@@ -65,14 +65,14 @@ class ConstraintTest(absltest.TestCase):
     pos = jp.ones(3)
 
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.REFSAFE
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     k, *_ = constraint._kbi(mx, solimp, solref, pos)
     self.assertEqual(k, 1 / (0.99**2 * timeconst**2))
 
   def test_disable_constraint(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT
-    ne, nf, nl, nc = mjx.count_constraints(m)
+    ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 0)
     self.assertEqual(nf, 0)
     self.assertEqual(nl, 0)
@@ -83,7 +83,7 @@ class ConstraintTest(absltest.TestCase):
   def test_disable_equality(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY
-    ne, nf, nl, nc = mjx.count_constraints(m)
+    ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 0)
     self.assertEqual(nf, 0)
     self.assertEqual(nl, 2)
@@ -94,7 +94,7 @@ class ConstraintTest(absltest.TestCase):
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
-    ne, nf, nl, nc = mjx.count_constraints(m)
+    ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 10)
     self.assertEqual(nf, 0)
     self.assertEqual(nl, 2)

mjx/mujoco/mjx/_src/device.py

@@ -1,319 +0,0 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# 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.
-# ==============================================================================
-"""Get and put mujoco data on/off device."""
-
-import copy
-import dataclasses
-from typing import Any, Dict, Iterable, List, Union, overload
-import warnings
-
-import jax
-from jax import numpy as jp
-import mujoco
-from mujoco.mjx._src import collision_driver
-from mujoco.mjx._src import mesh
-from mujoco.mjx._src import types
-import numpy as np
-
-_MJ_TYPE_ATTR = {
-    mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,),
-    mujoco.mjtDyn: (mujoco.MjModel.actuator_dyntype,),
-    mujoco.mjtEq: (mujoco.MjModel.eq_type,),
-    mujoco.mjtGain: (mujoco.MjModel.actuator_gaintype,),
-    mujoco.mjtTrn: (mujoco.MjModel.actuator_trntype,),
-    mujoco.mjtCone: (
-        mujoco.MjModel.opt,
-        mujoco.MjOption.cone,
-    ),
-    mujoco.mjtIntegrator: (
-        mujoco.MjModel.opt,
-        mujoco.MjOption.integrator,
-    ),
-    mujoco.mjtSolver: (
-        mujoco.MjModel.opt,
-        mujoco.MjOption.solver,
-    ),
-}
-
-_TYPE_MAP = {
-    mujoco._structs._MjContactList: types.Contact,  # pylint: disable=protected-access
-    mujoco.MjData: types.Data,
-    mujoco.MjModel: types.Model,
-    mujoco.MjOption: types.Option,
-    mujoco.MjStatistic: types.Statistic,
-    mujoco.mjtBias: types.BiasType,
-    mujoco.mjtCone: types.ConeType,
-    mujoco.mjtDisableBit: types.DisableBit,
-    mujoco.mjtDyn: types.DynType,
-    mujoco.mjtEq: types.EqType,
-    mujoco.mjtGain: types.GainType,
-    mujoco.mjtIntegrator: types.IntegratorType,
-    mujoco.mjtSolver: types.SolverType,
-    mujoco.mjtTrn: types.TrnType,
-}
-
-_TRANSFORMS = {
-    (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Data, 'cam_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Model, 'cam_mat0'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Contact, 'frame'): (
-        lambda x: x.reshape(x.shape[:-1] + (3, 3))  # pylint: disable=g-long-lambda
-        if x is not None and x.shape[0]
-        else jp.zeros((0, 3, 3))
-    ),
-}
-
-_INVERSE_TRANSFORMS = {
-    (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
-    (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
-    (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
-    (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
-    (types.Data, 'cam_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
-    (types.Model, 'cam_mat0'): lambda x: x.reshape(x.shape[:-2] + (9,)),
-    (types.Contact, 'frame'): (
-        lambda x: x.reshape(x.shape[:-2] + (9,))  # pylint: disable=g-long-lambda
-        if x is not None and x.shape[0]
-        else jp.zeros((0, 9))
-    ),
-}
-
-_DERIVED = mesh.DERIVED.union(
-    # efc_J is dense in MJX, sparse in MJ. ignore for now.
-    {(types.Data, 'efc_J'), (types.Option, 'has_fluid_params')}
-)
-
-
-def _model_derived(value: mujoco.MjModel) -> Dict[str, Any]:
-  return {k: jax.device_put(v) for k, v in mesh.get(value).items()}
-
-
-def _data_derived(value: mujoco.MjData) -> Dict[str, Any]:
-  return {'efc_J': jax.device_put(value.efc_J)}
-
-
-def _option_derived(value: types.Option) -> Dict[str, Any]:
-  has_fluid = (
-      value.density > 0 or value.viscosity > 0 or (value.wind != 0.0).any()
-  )
-  return {'has_fluid_params': has_fluid}
-
-
-def _validate(m: mujoco.MjModel):
-  """Validates that an mjModel is compatible with MJX."""
-
-  # check enum types
-  for mj_type, attrs in _MJ_TYPE_ATTR.items():
-    val = m
-    for attr in attrs:
-      val = attr.fget(val)  # pytype: disable=attribute-error
-
-    typs = set(val) if isinstance(val, Iterable) else {val}
-    unsupported_typs = typs - set(_TYPE_MAP[mj_type])
-    unsupported = [mj_type(t) for t in unsupported_typs]  # pylint: disable=too-many-function-args
-    if unsupported:
-      raise NotImplementedError(f'{unsupported} not implemented.')
-
-  if m.ntendon:
-    raise NotImplementedError('Tendons are not supported.')
-
-  # check condim
-  if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
-    raise NotImplementedError('Only condim=3 is supported.')
-
-  if m.body_gravcomp.any():
-    raise NotImplementedError('gravcomp is not supported')
-
-  # check collision geom types
-  for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
-    g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
-    if g1 == mujoco.mjtGeom.mjGEOM_PLANE and g2 in (
-        mujoco.mjtGeom.mjGEOM_PLANE,
-        mujoco.mjtGeom.mjGEOM_HFIELD,
-    ):
-      # MuJoCo does not collide planes with other planes or hfields
-      continue
-    if collision_driver.get_collision_fn((g1, g2)) is None:
-      raise NotImplementedError(f'({g1}, {g2}) collisions not implemented.')
-    *_, params = collision_driver.get_params(m, c)
-    margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0)
-    if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap:
-      raise NotImplementedError(
-          f'Margin and gap not implemented for ({g1}, {g2})'
-      )
-
-  # TODO(erikfrey): warn for high solver iterations, nefc, etc.
-
-  # mjNDISABLE is not a DisableBit flag, so must be explicitly ignored
-  disablebit_members = set(mujoco.mjtDisableBit.__members__.values()) - {
-      mujoco.mjtDisableBit.mjNDISABLE}
-  unsupported_disable = disablebit_members - {
-      mujoco.mjtDisableBit(t.value) for t in types.DisableBit
-  }
-  for f in unsupported_disable:
-    if f & m.opt.disableflags:
-      warnings.warn(f'Ignoring disable flag {f.name}.')
-
-  # mjNENABLE is not an EnableBit flag, so must be explicitly ignored
-  unsupported_enable = set(mujoco.mjtEnableBit.__members__.values()) - {
-      mujoco.mjtEnableBit.mjNENABLE
-  }
-  for f in unsupported_enable:
-    if f & m.opt.enableflags:
-      warnings.warn(f'Ignoring enable flag {f.name}.')
-
-  if not np.allclose(m.dof_frictionloss, 0):
-    raise NotImplementedError('dof_frictionloss is not implemented.')
-
-
-@overload
-def device_put(value: mujoco.MjData) -> types.Data:
-  ...
-
-
-@overload
-def device_put(value: mujoco.MjModel) -> types.Model:
-  ...
-
-
-def device_put(value):
-  """Places mujoco data onto a device.
-
-  Args:
-    value: a mujoco struct to transfer
-
-  Returns:
-    on-device MJX struct reflecting the input value
-  """
-  warnings.warn(
-      'device_put is deprecated, use put_model and put_data instead',
-      category=DeprecationWarning,
-  )
-
-  clz = _TYPE_MAP.get(type(value))
-  if clz is None:
-    raise NotImplementedError(f'{type(value)} is not supported for device_put.')
-
-  if isinstance(value, mujoco.MjModel):
-    _validate(value)  # type: ignore
-
-  init_kwargs = {}
-  for f in dataclasses.fields(clz):  # type: ignore
-    if (clz, f.name) in _DERIVED:
-      continue
-
-    field_value = getattr(value, f.name)
-    if (clz, f.name) in _TRANSFORMS:
-      field_value = _TRANSFORMS[(clz, f.name)](field_value)
-
-    if f.type is jax.Array:
-      field_value = jax.device_put(field_value)
-    elif type(field_value) in _TYPE_MAP.keys():
-      field_value = device_put(field_value)
-
-    init_kwargs[f.name] = copy.copy(field_value)
-
-  derived_kwargs = {}
-  if isinstance(value, mujoco.MjModel):
-    derived_kwargs = _model_derived(value)
-  elif isinstance(value, mujoco.MjData):
-    derived_kwargs = _data_derived(value)
-  elif isinstance(value, mujoco.MjOption):
-    derived_kwargs = _option_derived(value)
-
-  return clz(**init_kwargs, **derived_kwargs)  # type: ignore
-
-
-@overload
-def device_get_into(
-    result: Union[mujoco.MjData, List[mujoco.MjData]], value: types.Data
-):
-  ...
-
-
-def device_get_into(result, value):
-  """Transfers data off device into a mujoco MjData.
-
-  Data on device often has a batch dimension which adds (N,) to the beginning
-  of each array shape where N = batch size.
-
-  If result is a single MjData, arrays are copied over with the batch dimension
-  intact. If result is a list, the list must be length N and will be populated
-  with distinct MjData structs where the batch dimension is stripped.
-
-  Args:
-    result: struct (or list of structs) to transfer into
-    value: device value to transfer
-
-  Raises:
-    RuntimeError: if result length doesn't match data batch size
-  """
-  warnings.warn(
-      'device_get_into is deprecated, use get_data instead',
-      category=DeprecationWarning,
-  )
-
-  value = jax.device_get(value)
-
-  if isinstance(result, list):
-    array_shapes = [s.shape for s in jax.tree_util.tree_flatten(value)[0]]
-
-    if any(len(s) < 1 or s[0] != array_shapes[0][0] for s in array_shapes):
-      raise ValueError('unrecognizable batch dimension in value')
-
-    batch_size = array_shapes[0][0]
-
-    if len(result) != batch_size:
-      raise ValueError(
-          f"result length ({len(result)}) doesn't match value batch size"
-          f' ({batch_size})'
-      )
-
-    for i in range(batch_size):
-      value_i = jax.tree_map(lambda x, i=i: x[i], value)
-      device_get_into(result[i], value_i)
-
-  else:
-    if isinstance(result, mujoco.MjData):
-      ncon = value.contact.dist.shape[0]
-      nefc = value.efc_J.shape[0]
-      mujoco._functions._realloc_con_efc(  # pylint: disable=protected-access
-          result, ncon=ncon, nefc=nefc
-      )
-      result.ncon = ncon
-      result.nefc = nefc
-      efc_start = nefc - ncon * 4
-      result.contact.efc_address[:] = np.arange(efc_start, nefc, 4)
-      result.contact.dim[:] = 3
-
-    for f in dataclasses.fields(value):  # type: ignore
-      if (type(value), f.name) in _DERIVED:
-        continue
-
-      field_value = getattr(value, f.name)
-
-      if (type(value), f.name) in _INVERSE_TRANSFORMS:
-        field_value = _INVERSE_TRANSFORMS[(type(value), f.name)](field_value)
-
-      if type(field_value) in _TYPE_MAP.values():
-        device_get_into(getattr(result, f.name), field_value)
-        continue
-
-      try:
-        setattr(result, f.name, field_value)
-      except AttributeError:
-        getattr(result, f.name)[:] = field_value

mjx/mujoco/mjx/_src/device_test.py

@@ -1,203 +0,0 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# 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 moving mujoco structs on and off device."""
-
-import dataclasses
-
-from absl.testing import absltest
-from absl.testing import parameterized
-import jax
-from jax import numpy as jp
-import mujoco
-from mujoco import mjx
-from mujoco.mjx._src import device
-from mujoco.mjx._src import test_util
-from mujoco.mjx._src import types
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.dataclasses import PyTreeNode
-# pylint: enable=g-importing-member
-import numpy as np
-
-
-def _assert_eq(testcase, a, b, attr=None, name=None):
-  if (type(a), attr) in device._DERIVED:
-    return
-
-  if attr:
-    a, b = getattr(a, attr), getattr(b, attr)
-
-  if isinstance(a, PyTreeNode):
-    for field in dataclasses.fields(a):
-      _assert_eq(testcase, a, b, field.name, type(a).__name__)
-    return
-
-  typ = {'Model': types.Model, 'Data': types.Data,
-         'Contact': types.Contact}.get(name)
-  if (typ, attr) in device._TRANSFORMS:
-    b = device._TRANSFORMS[(typ, attr)](b)
-
-  err_msg = f'mismatch: {attr} in {name}'
-  if not hasattr(b, 'shape') or not b.shape:
-    testcase.assertEqual(a, b, err_msg)
-    return
-
-  a, b = np.array(a), np.array(b)
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=1e-8)
-
-
-class DeviceTest(parameterized.TestCase):
-
-  @parameterized.parameters('constraints.xml', 'pendula.xml')
-  def testdevice_put(self, fname):
-    """Test putting MjData and MjModel on device."""
-    m = test_util.load_test_file(fname)
-    # advance state to ensure non-zero fields
-    d = mujoco.MjData(m)
-    for _ in range(10):
-      mujoco.mj_step(m, d)
-
-    _assert_eq(self, mjx.device_put(d), d)
-    _assert_eq(self, mjx.device_put(m), m)
-
-  @parameterized.parameters('constraints.xml', 'pendula.xml')
-  def testdevice_get(self, fname):
-    """Test getting MjData from a device."""
-    m = test_util.load_test_file(fname)
-    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
-    mx = device.device_put(m)
-    dx = mjx.make_data(mx)
-    d = mujoco.MjData(m)
-    device.device_get_into(d, dx)
-    _assert_eq(self, dx, d)
-
-  @parameterized.parameters('constraints.xml', 'pendula.xml')
-  def testdevice_get_batched(self, fname):
-    """Test getting MjData from a device."""
-    m = test_util.load_test_file(fname)
-    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
-    mx = device.device_put(m)
-    batch_size = 32
-
-    # create mjx_data and batch it
-    dx = mjx.make_data(mx)
-    dx = jax.tree_map(
-        lambda x: jp.repeat(x, batch_size).reshape((batch_size,) + x.shape),
-        dx,
-    )
-    ds = [mujoco.MjData(m) for _ in range(batch_size - 1)]
-
-    with self.assertRaises(ValueError):
-      device.device_get_into(ds, dx)
-
-    ds = [mujoco.MjData(m) for _ in range(batch_size)]
-    device.device_get_into(ds, dx)
-    dx = jax.device_get(dx)  # faster indexing for testing
-    for i in range(batch_size):
-      _assert_eq(self, jax.tree_map(lambda x, i=i: x[i], dx), ds[i])
-
-
-class ValidateInputTest(absltest.TestCase):
-
-  def test_solver(self):
-    m = mujoco.MjModel.from_xml_string(
-        '<mujoco><option solver="PGS"/><worldbody/></mujoco>'
-    )
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_integrator(self):
-    m = mujoco.MjModel.from_xml_string(
-        '<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
-    )
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_cone(self):
-    m = mujoco.MjModel.from_xml_string(
-        '<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
-    )
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_dyn(self):
-    m = test_util.load_test_file('pendula.xml')
-    m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_gain(self):
-    m = test_util.load_test_file('pendula.xml')
-    m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_bias(self):
-    m = test_util.load_test_file('pendula.xml')
-    m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_condim(self):
-    m = test_util.load_test_file('constraints.xml')
-    for i in [1, 4, 6]:
-      m.geom_condim[0] = i
-      with self.assertRaises(NotImplementedError):
-        mjx.device_put(m)
-
-  def test_geoms(self):
-    m = mujoco.MjModel.from_xml_string("""
-      <mujoco>
-        <worldbody>
-          <body>
-            <joint axis="1 0 0" type="free"/>
-            <geom size="0.2 0.2 0.2" type="box"/>
-          </body>
-          <body>
-            <joint axis="1 0 0" type="free"/>
-            <geom size="0.1 0.1" type="cylinder"/>
-          </body>
-        </worldbody>
-      </mujoco>
-    """)
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-  def test_tendon(self):
-    m = mujoco.MjModel.from_xml_string("""
-      <mujoco>
-        <worldbody>
-        <body name="left_thigh" pos="0 0.1 -0.04">
-          <joint axis="0 1 0" name="left_hip_y" type="hinge"/>
-          <geom fromto="0 0 0 0 -0.01 -.34" name="left_thigh1" size="0.06" type="capsule"/>
-          <body name="left_shin" pos="0 -0.01 -0.403">
-            <joint axis="0 -1 0" name="left_knee" pos="0 0 .02" range="-160 -2" type="hinge"/>
-            <geom fromto="0 0 0 0 0 -.3" name="left_shin1" size="0.049" type="capsule"/>
-          </body>
-        </body>
-        </worldbody>
-        <tendon>
-          <fixed name="left_hipknee">
-            <joint coef="-1" joint="left_hip_y"/>
-            <joint coef="1" joint="left_knee"/>
-          </fixed>
-        </tendon>
-      </mujoco>
-    """)
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
-
-if __name__ == '__main__':
-  absltest.main()

mjx/mujoco/mjx/_src/forward.py

@@ -310,7 +310,7 @@ def rungekutta4(m: Model, d: Data) -> Data:
 
   kqvel = d.qvel  # intermediate RK solution
   # RK solutions sum
-  qvel, qacc, act_dot = jax.tree_map(
+  qvel, qacc, act_dot = jax.tree_util.tree_map(
       lambda k: B[0] * k, (kqvel, d.qacc, d.act_dot)
   )
   integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep)
@@ -318,7 +318,7 @@ def rungekutta4(m: Model, d: Data) -> Data:
   def f(carry, x):
     qvel, qacc, act_dot, kqvel, d = carry
     a, b, t = x  # tableau numbers
-    dqvel, dqacc, dact_dot = jax.tree_map(
+    dqvel, dqacc, dact_dot = jax.tree_util.tree_map(
         lambda k: a * k, (kqvel, d.qacc, d.act_dot)
     )
     # get intermediate RK solutions

mjx/mujoco/mjx/_src/io.py

@@ -22,15 +22,14 @@ from jax import numpy as jp
 import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
-from mujoco.mjx._src import mesh
 from mujoco.mjx._src import support
 from mujoco.mjx._src import types
 import numpy as np
 import scipy
 
 
-def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
+def _make_option(o: mujoco.MjOption) -> types.Option:
-  """Puts mujoco.MjOption onto a device, resulting in mjx.Option."""
+  """Returns mjx.Option given mujoco.MjOption."""
   if o.integrator not in set(types.IntegratorType):
     raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
 
@@ -47,38 +46,20 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
     if o.enableflags & 2**i:
       raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
 
-  static_fields = {
+  fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()}
-      f.name: copy.copy(getattr(o, f.name))
+  fields['integrator'] = types.IntegratorType(o.integrator)
-      for f in types.Option.fields()
+  fields['cone'] = types.ConeType(o.cone)
-      if f.type in (int, bytes, np.ndarray)
+  fields['jacobian'] = types.JacobianType(o.jacobian)
-  }
+  fields['solver'] = types.SolverType(o.solver)
-  static_fields['integrator'] = types.IntegratorType(o.integrator)
+  fields['disableflags'] = types.DisableBit(o.disableflags)
-  static_fields['cone'] = types.ConeType(o.cone)
+  fields['has_fluid_params'] = o.density > 0 or o.viscosity > 0 or o.wind.any()
-  static_fields['jacobian'] = types.JacobianType(o.jacobian)
-  static_fields['solver'] = types.SolverType(o.solver)
-  static_fields['disableflags'] = types.DisableBit(o.disableflags)
 
-  device_fields = {
+  return types.Option(**fields)
-      f.name: copy.copy(getattr(o, f.name))
-      for f in types.Option.fields()
-      if f.type is jax.Array
-  }
-  device_fields = jax.device_put(device_fields, device=device)
-
-  has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
-
-  return types.Option(
-      has_fluid_params=has_fluid_params,
-      **static_fields,
-      **device_fields,
-  )
 
 
-def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic:
+def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic:
   """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic."""
-  return types.Statistic(
+  return types.Statistic(meaninertia=s.meaninertia)
-      meaninertia=jax.device_put(s.meaninertia, device=device)
-  )
 
 
 def put_model(m: mujoco.MjModel, device=None) -> types.Model:
@@ -93,17 +74,21 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
   if m.body_gravcomp.any():
     raise NotImplementedError('gravcomp is not supported')
 
-  # check collision geom types
+  for g1, g2, ip in collision_driver.geom_pairs(m):
-  for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
+    t1, t2 = m.geom_type[[g1, g2]]
-    g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
+    # check collision function exists for type pair
-    if collision_driver.get_collision_fn((g1, g2)) is None:
+    if not collision_driver.has_collision_fn(t1, t2):
-      raise NotImplementedError(f'({g1}, {g2}) has no collision function')
+      t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
-    *_, params = collision_driver.get_params(m, c)
+      raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.')
-    margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0)
+    # margin/gap not supported for geoms
-    if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap:
+    if mujoco.mjtGeom.mjGEOM_MESH in (t1, t2):
-      raise NotImplementedError(
+      if ip != -1:
-          f'Margin and gap not implemented for ({g1}, {g2})'
+        margin = m.pair_margin[ip]
-      )
+      else:
+        margin = m.geom_margin[g1] + m.geom_margin[g2]
+      if margin.any():
+        t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
+        raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.')
 
   for enum_field, enum_type, mj_type in (
       (m.actuator_biastype, types.BiasType, mujoco.mjtBias),
@@ -121,40 +106,24 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
   if not np.allclose(m.dof_frictionloss, 0):
     raise NotImplementedError('dof_frictionloss is not implemented.')
 
-  opt = _put_option(m.opt, device=device)
+  fields = {f.name: getattr(m, f.name) for f in types.Model.fields()}
-  stat = _put_statistic(m.stat, device=device)
+  fields['geom_rgba'] = fields['geom_rgba'].reshape((-1, 4))
+  fields['mat_rgba'] = fields['mat_rgba'].reshape((-1, 4))
+  fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
+  fields['opt'] = _make_option(m.opt)
+  fields['stat'] = _make_statistic(m.stat)
+  model = types.Model(**{k: copy.copy(v) for k, v in fields.items()})
 
-  static_fields = {
+  return jax.device_put(model, device=device)
-      f.name: getattr(m, f.name)
-      for f in types.Model.fields()
-      if f.type in (int, bytes, np.ndarray)
-  }
-  static_fields['geom_rgba'] = static_fields['geom_rgba'].reshape((-1, 4))
-  static_fields['mat_rgba'] = static_fields['mat_rgba'].reshape((-1, 4))
-
-  device_fields = {
-      f.name: copy.copy(getattr(m, f.name))  # copy because device_put is async
-      for f in types.Model.fields()
-      if f.type is jax.Array
-  }
-  device_fields['cam_mat0'] = device_fields['cam_mat0'].reshape((-1, 3, 3))
-  device_fields.update(mesh.get(m))
-  device_fields = jax.device_put(device_fields, device=device)
-
-  return types.Model(
-      opt=opt,
-      stat=stat,
-      **static_fields,
-      **device_fields,
-  )
 
 
 def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   """Allocate and initialize Data."""
-
+  dim = collision_driver.make_condim(m)
-  ncon = collision_driver.ncon(m)
+  efc_type = constraint.make_efc_type(m, dim)
-  ne, nf, nl, nc = constraint.count_constraints(m)
+  efc_address = constraint.make_efc_address(efc_type, dim)
-  nefc = ne + nf + nl + nc
+  ne, nf, nl, nc = constraint.counts(efc_type)
+  ncon, nefc = dim.size, ne + nf + nl + nc
 
   zero_0 = jp.zeros(0, dtype=float)
   zero_nv = jp.zeros(m.nv, dtype=float)
@@ -170,8 +139,28 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=float)
   zero_nm = jp.zeros(m.nM, dtype=float)
 
-  # create first d to get num contacts and nc
+  contact = types.Contact(
+      dist=jp.zeros(ncon),
+      pos=jp.zeros((ncon, 3)),
+      frame=jp.zeros((ncon, 3, 3)),
+      includemargin=jp.zeros(ncon),
+      friction=jp.zeros((ncon, 5)),
+      solref=jp.zeros((ncon, mujoco.mjNREF)),
+      solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
+      solimp=jp.zeros((ncon, mujoco.mjNIMP)),
+      dim=dim,
+      geom1=jp.zeros(ncon, dtype=int) - 1,
+      geom2=jp.zeros(ncon, dtype=int) - 1,
+      geom=jp.zeros((ncon, 2), dtype=int) - 1,
+      efc_address=efc_address,
+  )
+
   d = types.Data(
+      ne=ne,
+      nf=nf,
+      nl=nl,
+      nefc=nefc,
+      ncon=ncon,
       solver_niter=jp.array(0, dtype=int),
       time=jp.array(0.0),
       qpos=jp.array(m.qpos0),
@@ -206,7 +195,8 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
       qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
       qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
-      contact=types.Contact.zero(ncon),
+      contact=contact,
+      efc_type=efc_type,
       efc_J=jp.zeros((nefc, m.nv), dtype=float),
       efc_frictionloss=zero_nefc,
       efc_D=zero_nefc,
@@ -228,11 +218,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   return d
 
 
-def _get_contact(
+def _get_contact(c: mujoco._structs._MjContactList, cx: types.Contact):
-    c: mujoco._structs._MjContactList,
-    cx: types.Contact,
-    efc_start: int,
-):
   """Converts mjx.Contact to mujoco._structs._MjContactList."""
   con_id = np.nonzero(cx.dist <= 0)[0]
   for field in types.Contact.fields():
@@ -241,10 +227,6 @@ def _get_contact(
       value = value.reshape((-1, 9))
     getattr(c, field.name)[:] = value
 
-  ncon = cx.dist.shape[0]
-  c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
-  c.dim[:] = 3
-
 
 def get_data(
     m: mujoco.MjModel, d: types.Data
@@ -278,13 +260,6 @@ def get_data_into(
   d = jax.device_get(d)
 
   batch_size = d.qpos.shape[0] if batched else 1
-  ne, nf, nl, nc = constraint.count_constraints(m, d)
-  efc_type = np.array([
-      mujoco.mjtConstraint.mjCNSTR_EQUALITY,
-      mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF,
-      mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT,
-      mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
-  ]).repeat([ne, nf, nl, nc])
 
   dof_i, dof_j = [], []
   for i in range(m.nv):
@@ -295,12 +270,11 @@ def get_data_into(
       j = m.dof_parentid[j]
 
   for i in range(batch_size):
-    d_i = jax.tree_map(lambda x, i=i: x[i], d) if batched else d
+    d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
     result_i = result[i] if batched else result
     ncon = (d_i.contact.dist <= 0).sum()
     efc_active = (d_i.efc_J != 0).any(axis=1)
-    efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
+    nefc = int(efc_active.sum())
-    nefc, nc = int(efc_active.sum()), int((efc_active & efc_con).sum())
     result_i.nnzJ = nefc * m.nv
     if ncon != result_i.ncon or nefc != result_i.nefc:
       mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc)  # pylint: disable=protected-access
@@ -310,61 +284,63 @@ def get_data_into(
 
     for field in types.Data.fields():
       if field.name == 'contact':
-        _get_contact(result_i.contact, d_i.contact, nefc - nc)
+        _get_contact(result_i.contact, d_i.contact)
+        # efc_address must be updated because rows were deleted above:
+        efc_map = np.cumsum(efc_active) - 1
+        result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address]
         continue
 
       value = getattr(d_i, field.name)
 
-      if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
+      if field.name in ('nefc', 'ncon'):
+        value = {'nefc': nefc, 'ncon': ncon}[field.name]
+      elif field.name.endswith('xmat') or field.name == 'ximat':
         value = value.reshape((-1, 9))
-
+      elif field.name.startswith('efc_'):
-      if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
         value = value[efc_active]
-
+        if field.name == 'efc_J':
-      if field.name == 'efc_J':
+          value = value.reshape(-1)
-        value = value[efc_active].reshape(-1)
+      elif field.name == 'qM' and not support.is_sparse(m):
-
-      if field.name == 'qM' and not support.is_sparse(m):
         value = value[dof_i, dof_j]
-
+      elif field.name == 'qLD' and not support.is_sparse(m):
-      if field.name == 'qLD' and not support.is_sparse(m):
         value = value[dof_i, dof_j]
-
+      elif field.name == 'qLDiagInv' and not support.is_sparse(m):
-      if field.name == 'qLDiagInv' and not support.is_sparse(m):
         value = np.ones(m.nv)
 
-      if value.shape:
+      if isinstance(value, np.ndarray) and value.shape:
         getattr(result_i, field.name)[:] = value
       else:
         setattr(result_i, field.name, value)
 
-    result_i.efc_type[:] = efc_type[efc_active]
 
-
+def _make_contact(
-def _put_contact(
+    c: mujoco._structs._MjContactList,
-    c: mujoco._structs._MjContactList, ncon: int, device=None
+    dim: np.ndarray,
+    efc_address: np.ndarray,
 ) -> types.Contact:
-  """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact."""
+  """Converts mujoco.structs._MjContactList into mjx.Contact."""
-  fields = {
+  fields = {f.name: getattr(c, f.name) for f in types.Contact.fields()}
-      f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields()
-  }
   fields['frame'] = fields['frame'].reshape((-1, 3, 3))
-  pad_size = ncon - c.dist.shape[0]
+  pad_size = dim.size - c.dist.shape[0]
   pad_fn = lambda x: np.concatenate(
       (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
   )
-  fields = jax.tree_map(pad_fn, fields)
+  fields = jax.tree_util.tree_map(pad_fn, fields)
   fields['dist'][-pad_size:] = np.inf
-  fields = jax.device_put(fields, device=device)
+  # TODO(erikfrey): move contacts to appropriate dim index
+  fields['dim'] = dim
+  fields['efc_address'] = efc_address
 
   return types.Contact(**fields)
 
 
 def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
   """Puts mujoco.MjData onto a device, resulting in mjx.Data."""
-  ncon = collision_driver.ncon(m)
+  dim = collision_driver.make_condim(m)
-  ne, nf, nl, nc = constraint.count_constraints(m)
+  efc_type = constraint.make_efc_type(m, dim)
-  nefc = ne + nf + nl + nc
+  efc_address = constraint.make_efc_address(efc_type, dim)
+  ne, nf, nl, nc = constraint.counts(efc_type)
+  ncon, nefc = dim.size, ne + nf + nl + nc
 
   for d_val, val, name in (
       (d.ncon, ncon, 'ncon'),
@@ -376,12 +352,9 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
     if d_val > val:
       raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
 
-  fields = {
+  fields = {f.name: getattr(d, f.name) for f in types.Data.fields()}
-      f.name: copy.copy(getattr(d, f.name))  # copy because device_put is async
-      for f in types.Data.fields()
-      if f.type is jax.Array
-  }
 
+  # MJX prefers square matrices for these fields:
   for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
     fields[fname] = fields[fname].reshape((-1, 3, 3))
 
@@ -409,7 +382,7 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
       value_beg = sum([ne, nf, nl][:i])
       d_beg = sum([d.ne, d.nf, d.nl][:i])
       size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i]
-      value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size]
+      value[value_beg : value_beg + size] = fields[fname][d_beg : d_beg + size]
     fields[fname] = value
 
   # convert qM and qLD if jacobian is dense
@@ -424,7 +397,12 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
       fields['qLD'] = np.zeros((m.nv, m.nv))
     fields['qLDiagInv'] = np.zeros(0)
 
-  fields = jax.device_put(fields, device=device)
+  fields['contact'] = _make_contact(d.contact, dim, efc_address)
-  fields['contact'] = _put_contact(d.contact, ncon, device=device)
+  fields.update(
+      dict(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)
+  )
 
-  return types.Data(**fields)
+  # copy because device_put is async:
+  data = types.Data(**{k: copy.copy(v) for k, v in fields.items()})
+
+  return jax.device_put(data, device=device)

mjx/mujoco/mjx/_src/io_test.py

@@ -108,10 +108,6 @@ class ModelIOTest(parameterized.TestCase):
     np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)
     np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref)
     np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos)
-    self.assertLen(mx.geom_convex_face, 6)
-    self.assertLen(mx.geom_convex_vert, 6)
-    self.assertLen(mx.geom_convex_edge_dir, 6)
-    self.assertLen(mx.geom_convex_facenormal, 6)
 
     np.testing.assert_allclose(mx.jnt_type, m.jnt_type)
     np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr)
@@ -256,6 +252,7 @@ class DataIOTest(parameterized.TestCase):
     nv = 19
     nefc = 185
 
+    self.assertEqual(d.nefc, nefc)
     self.assertEqual(d.qpos.shape, (nq,))
     self.assertEqual(d.qvel.shape, (nv,))
     self.assertEqual(d.act.shape, (0,))
@@ -423,8 +420,6 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(d_2.efc_J, d.efc_J)
     self.assertEqual(d_2.efc_aref.shape, (8,))  # nefc
     np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
-
-    # efc_address is created on demand
     np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
 
   def test_get_data_batched(self):
@@ -435,7 +430,7 @@ class DataIOTest(parameterized.TestCase):
     mujoco.mj_step(m, d, 2)
     dx = mjx.put_data(m, d)
     # second data in batch has contact dist > 0, disables contact
-    dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
+    dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
     ds = mjx.get_data(m, dx_b)
     self.assertLen(ds, 2)
     np.testing.assert_allclose(ds[0].qpos, d.qpos)

mjx/mujoco/mjx/_src/mesh.py

@@ -15,14 +15,15 @@
 """Mesh processing."""
 
 import collections
-import dataclasses
 import itertools
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Tuple
 import warnings
 
-import mujoco
+import jax
+from jax import numpy as jp
 # pylint: disable=g-importing-member
-from mujoco.mjx._src.types import GeomType
+from mujoco.mjx._src.collision_types import ConvexInfo
+from mujoco.mjx._src.collision_types import GeomInfo
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 import numpy as np
@@ -30,37 +31,7 @@ from scipy import spatial
 import trimesh
 
 
-_BOX_CORNERS = list(itertools.product((-1, 1), (-1, 1), (-1, 1)))
-# pyformat: disable
-# Rectangular box faces using a counter-clockwise winding order convention.
-_BOX_FACES = [
-    0, 4, 5, 1,  # left
-    0, 2, 6, 4,  # bottom
-    6, 7, 5, 4,  # front
-    2, 3, 7, 6,  # right
-    1, 5, 7, 3,  # top
-    0, 1, 3, 2,  # back
-]
-# pyformat: enable
 _MAX_HULL_FACE_VERTICES = 20
-_CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {}
-_DERIVED_ARGS = [
-    'geom_convex_face',
-    'geom_convex_vert',
-    'geom_convex_edge_dir',
-    'geom_convex_facenormal',
-    'geom_convex_edge',
-    'geom_convex_edge_face_normal',
-]
-DERIVED = {(Model, d) for d in _DERIVED_ARGS}
-
-
-def _box(size: np.ndarray):
-  """Creates a mesh for a box with rectangular faces."""
-  box_corners = np.array(_BOX_CORNERS)
-  vert = box_corners * size.reshape(-1, 3)
-  face = np.array([_BOX_FACES]).reshape(-1, 4)
-  return vert, face
 
 
 def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
@@ -170,16 +141,7 @@ def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray:
   return hull_point_idx
 
 
-@dataclasses.dataclass
+def _merge_coplanar(m: Model, tm: trimesh.Trimesh, meshid: int) -> np.ndarray:
-class MeshInfo:
-  name: str
-  vert: np.ndarray
-  face: np.ndarray
-  convex_vert: Optional[np.ndarray]
-  convex_face: Optional[np.ndarray]
-
-
-def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
   """Merges coplanar facets."""
   if not tm.facets:
     return tm.faces.copy()  # no facets
@@ -204,9 +166,11 @@ def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
 
     # resize faces that exceed max polygon vertices
     if face.shape[0] > _MAX_HULL_FACE_VERTICES:
+      name = m.names[m.name_meshadr[meshid]:]
+      name = name[:name.find(b'\x00')].decode('utf-8')
       warnings.warn(
-          f'Mesh "{mesh_info.name}" has a coplanar face with more than'
+          f'Mesh "{name}" has a coplanar face with more than '
-          f' {_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance '
+          f'{_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance '
           'issues and inaccuracies in collision detection. Consider '
           'decimating the mesh.'
       )
@@ -231,106 +195,99 @@ def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
   return np.concatenate([faces, facets])
 
 
-def _mesh_info(
+def box(info: GeomInfo) -> ConvexInfo:
-    m: mujoco.MjModel,
+  """Creates a box with rectangular faces."""
-) -> List[MeshInfo]:
+  vert = np.array(
-  """Extracts mesh info from MjModel."""
+      list(itertools.product((-1, 1), (-1, 1), (-1, 1))), dtype=float
-  mesh_infos = []
-  for i in range(m.nmesh):
-    name = mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_MESH.value, i)
-
-    last = (i + 1) >= m.nmesh
-    face_start = m.mesh_faceadr[i]
-    face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0]
-    face = m.mesh_face[face_start:face_end]
-
-    vert_start = m.mesh_vertadr[i]
-    vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0]
-    vert = m.mesh_vert[vert_start:vert_end]
-
-    graphadr = m.mesh_graphadr[i]
-    if graphadr < 0:
-      mesh_infos.append(MeshInfo(name, vert, face, None, None))
-      continue
-
-    graph = m.mesh_graph[graphadr:]
-    numvert, numface = graph[0], graph[1]
-
-    # unused vert_edgeadr
-    # vert_edgeadr = graph[2 : numvert + 2]
-    last_idx = numvert + 2
-
-    vert_globalid = graph[last_idx : last_idx + numvert]
-    last_idx += numvert
-
-    # unused edge_localid
-    # edge_localid = graph[last_idx : last_idx + numvert + 3 * numface]
-    last_idx += numvert + 3 * numface
-
-    face_globalid = graph[last_idx : last_idx + 3 * numface]
-    face_globalid = face_globalid.reshape((numface, 3))
-
-    convex_vert = vert[vert_globalid]
-    vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0])))
-    convex_face = np.vectorize(vertex_map.get)(face_globalid)
-    mesh_infos.append(MeshInfo(name, vert, face, convex_vert, convex_face))
-
-  return mesh_infos
-
-
-def _geom_mesh_kwargs(
-    mesh_info: MeshInfo,
-) -> Dict[str, np.ndarray]:
-  """Generates convex mesh attributes for mjx.Model."""
-  tm_convex = trimesh.Trimesh(
-      vertices=mesh_info.convex_vert, faces=mesh_info.convex_face
   )
+  # pyformat: disable
+  # rectangular box faces using a counter-clockwise winding order convention:
+  face = np.array(
+      [
+          0, 4, 5, 1,  # left
+          0, 2, 6, 4,  # bottom
+          6, 7, 5, 4,  # front
+          2, 3, 7, 6,  # right
+          1, 5, 7, 3,  # top
+          0, 1, 3, 2,  # back
+      ]
+  ).reshape((-1, 4))
+  # pyformat: enable
+  face_normal = _get_face_norm(vert, face)
+  edge, edge_face_normal = _get_edge_normals(face, face_normal)
+  edge_dir = _get_unique_edge_dir(vert, face)
+  face = vert[face]  # materialize full nface x nvert matrix
+
+  c = ConvexInfo(
+      info.pos,
+      info.mat,
+      vert,
+      face,
+      face_normal,
+      edge,
+      edge_face_normal,
+      edge_dir,
+  )
+  c = jax.tree_util.tree_map(jp.array, c)
+  vert = jax.vmap(jp.multiply, in_axes=(None, 0))(c.vert, info.size)
+  face = jax.vmap(jp.multiply, in_axes=(None, 0))(c.face, info.size)
+  c = c.replace(vert=vert, face=face)
+
+  return c
+
+
+def convex(m: Model, mesh_id: int, info: GeomInfo) -> ConvexInfo:
+  """Processes a mesh for use in convex collision algorithms.
+
+  Args:
+    m: an MJX model
+    mesh_id: the mesh id to process
+    info: pos, mat, size of this geom
+
+  Returns:
+    a convex mesh info
+  """
+  vert_beg = m.mesh_vertadr[mesh_id]
+  vert_end = m.mesh_vertadr[mesh_id + 1] if mesh_id < m.nmesh - 1 else None
+  vert = m.mesh_vert[vert_beg:vert_end]
+
+  graphadr = m.mesh_graphadr[mesh_id]
+  graph = m.mesh_graph[graphadr:]
+  graph_idx = 0
+
+  numvert, numface = graph[0], graph[1]
+  graph_idx += 2
+
+  # skip vert_edgeadr  (numvert,)
+  graph_idx += numvert
+  vert_globalid = graph[graph_idx : graph_idx + numvert]
+  graph_idx += numvert
+
+  # skip edge_localid  (numvert, 3)
+  graph_idx += numvert + 3 * numface
+  face_globalid = graph[graph_idx : graph_idx + 3 * numface].reshape((-1, 3))
+
+  vert = vert[vert_globalid]
+  vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0])))
+  face = np.vectorize(vertex_map.get)(face_globalid)
+
+  tm_convex = trimesh.Trimesh(vertices=vert, faces=face)
   vert = np.array(tm_convex.vertices)
-  face = _merge_coplanar(tm_convex, mesh_info)
+  face = _merge_coplanar(m, tm_convex, mesh_id)
-  facenormal = _get_face_norm(vert, face)
+  face_normal = _get_face_norm(vert, face)
-  edge, edge_face_normal = _get_edge_normals(face, facenormal)
+  edge, edge_face_normal = _get_edge_normals(face, face_normal)
-  return {
+  edge_dir = _get_unique_edge_dir(vert, face)
-      'geom_convex_face': vert[face],
+  face = vert[face]  # materialize full nface x nvert matrix
-      'geom_convex_face_vert_idx': face,
-      'geom_convex_vert': vert,
-      'geom_convex_edge_dir': _get_unique_edge_dir(vert, face),
-      'geom_convex_facenormal': facenormal,
-      'geom_convex_edge': edge,
-      'geom_convex_edge_face_normal': edge_face_normal,
-  }
 
+  c = ConvexInfo(
+      info.pos,
+      info.mat,
+      vert,
+      face,
+      face_normal,
+      edge,
+      edge_face_normal,
+      edge_dir,
+  )
 
-def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
+  return jax.tree_util.tree_map(jp.array, c)
-  """Derives geom mesh attributes for mjx.Model from MjModel."""
-  kwargs = {k: [] for k in _DERIVED_ARGS}
-  mesh_infos = _mesh_info(m)
-  geom_con = m.geom_conaffinity | m.geom_contype
-  for geomid in range(m.ngeom):
-    mesh_info = None
-    dataid = m.geom_dataid[geomid]
-    if not geom_con[geomid]:
-      # ignore visual-only meshes
-      kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
-      continue
-    elif m.geom_type[geomid] == GeomType.BOX:
-      vert, face = _box(m.geom_size[geomid])
-      mesh_info = MeshInfo(
-          name='box',
-          vert=vert,
-          face=face,
-          convex_vert=vert,
-          convex_face=face,
-      )
-    elif dataid < 0:
-      kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
-      continue
-
-    mesh_info = mesh_info or mesh_infos[dataid]
-    vert, face = mesh_info.vert, mesh_info.face
-    key = (hash(vert.data.tobytes()), hash(face.data.tobytes()))
-    if key not in _CONVEX_CACHE:
-      _CONVEX_CACHE[key] = _geom_mesh_kwargs(mesh_info)
-
-    kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS}
-
-  return kwargs

mjx/mujoco/mjx/_src/mesh_test.py

@@ -20,7 +20,7 @@ import numpy as np
 import trimesh
 
 
-class GeomMeshKwargsTest(absltest.TestCase):
+class MeshTest(absltest.TestCase):
 
   def test_pyramid(self):
     """Tests that a triangulated pyramid converts to merged coplanar faces."""
@@ -37,29 +37,21 @@ class GeomMeshKwargsTest(absltest.TestCase):
     tm = trimesh.Trimesh(vertices=vert, faces=face)
     tm_convex = trimesh.convex.convex_hull(tm)
     convex_vert = np.array(tm_convex.vertices)
-    convex_face = np.array(tm_convex.faces)
+    convex_face = mesh._merge_coplanar(None, tm_convex, 0)
-    mesh_info = mesh.MeshInfo(
-        name='test',
-        vert=vert,
-        face=face,
-        convex_vert=convex_vert,
-        convex_face=convex_face,
-    )
-    h = mesh._geom_mesh_kwargs(mesh_info)
 
     # get index of vertices in h['geom_convex_vert'] for vertices in vert
     dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile(
-        h['geom_convex_vert'], (vert.shape[0], 1)
+        convex_vert, (vert.shape[0], 1)
     )
     dist = (dist**2).sum(axis=1).reshape((vert.shape[0], -1))
     vidx = np.argmin(dist, axis=0)
 
     # check verts
-    np.testing.assert_array_equal(h['geom_convex_vert'], vert[vidx])
+    np.testing.assert_array_equal(convex_vert, vert[vidx])
 
     # check face vertices
     map_ = {v: k for k, v in enumerate(vidx)}
-    h_face = np.vectorize(map_.get)(h['geom_convex_face_vert_idx'])
+    h_face = np.vectorize(map_.get)(convex_face)
     face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()])
     expected_face_verts = sorted([
         (0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)])
@@ -69,7 +61,8 @@ class GeomMeshKwargsTest(absltest.TestCase):
     )
 
     # check edges
-    unique_edge = np.vectorize(map_.get)(h['geom_convex_edge_dir'])
+    edge_dir = mesh._get_unique_edge_dir(convex_vert, convex_face)
+    unique_edge = np.vectorize(map_.get)(edge_dir)
     unique_edge = np.array(sorted(unique_edge.tolist()))
     np.testing.assert_array_equal(
         unique_edge,
@@ -77,10 +70,11 @@ class GeomMeshKwargsTest(absltest.TestCase):
     )
 
     # face normals
-    self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3))
+    face_normal = mesh._get_face_norm(convex_vert, convex_face)
+    self.assertEqual(face_normal.shape, (5, 3))
 
     # face edges
-    edges = h['geom_convex_edge']
+    edges, edge_normal = mesh._get_edge_normals(convex_face, face_normal)
     edges = np.vectorize(map_.get)(edges)
     mask = edges[:, 0] != edges[:, 1]
     edges = edges[mask]
@@ -103,7 +97,6 @@ class GeomMeshKwargsTest(absltest.TestCase):
     )
 
     # face edge normals
-    edge_normal = h['geom_convex_edge_face_normal']
     edge_normal = edge_normal[mask]
     edge_normal = np.take_along_axis(
         edge_normal, sort_col_idx[..., None], axis=1

mjx/mujoco/mjx/_src/scan.py

@@ -62,7 +62,7 @@ def _take(obj: Y, idx: np.ndarray) -> Y:
       x = x.take(jp.array(idx), axis=0, mode='wrap')
     return x
 
-  return jax.tree_map(take, obj)
+  return jax.tree_util.tree_map(take, obj)
 
 
 def _q_bodyid(m: Model) -> np.ndarray:
@@ -120,7 +120,7 @@ def _nvmap(f: Callable[..., Y], *args) -> Y:
   args = [a if n is None else None for n, a in zip(np_args, args)]
 
   # remove empty args that we should not vmap over
-  args = jax.tree_map(lambda a: a if a.shape[0] else None, args)
+  args = jax.tree_util.tree_map(lambda a: a if a.shape[0] else None, args)
   in_axes = [None if a is None else 0 for a in args]
 
   def outer_f(*args, np_args=np_args):
@@ -322,7 +322,7 @@ def flat(
       [v if typ in flat_ else jp.concatenate(v) for v, typ in zip(y, out_types)]
       for y in ys
   ]
-  ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
+  ys = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *ys)
 
   # put concatenated results back in order
   reordered_ys = []
@@ -465,15 +465,15 @@ def body_tree(
         def index_sum(x, i=id_map, s=body_ids.size):
           return jax.ops.segment_sum(x, i, s)
 
-        y = jax.tree_map(index_sum, y)
+        y = jax.tree_util.tree_map(index_sum, y)
-        carry = y if carry is None else jax.tree_map(jp.add, carry, y)
+        carry = y if carry is None else jax.tree_util.tree_map(jp.add, carry, y)
     elif key in key_parents:
       ys = [key_y[p] for p in key_parents[key]]
-      y = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
+      y = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *ys)
       body_ids = np.concatenate([key_body_ids[p] for p in key_parents[key]])
       parent_ids = m.body_parentid[key_body_ids[key]]
       take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i)
-      carry = jax.tree_map(take_fn, y)
+      carry = jax.tree_util.tree_map(take_fn, y)
 
     f_args = [_take(arg, ids) for arg, ids in zip(args, key_in_take[key])]
     key_y[key] = _nvmap(f, carry, *f_args)
@@ -488,8 +488,8 @@ def body_tree(
     if len(out_types) > 1:
       y_typ = [y_[i] for y_ in y_typ]
     if typ != 'b':
-      y_typ = jax.tree_map(jp.concatenate, y_typ)
+      y_typ = jax.tree_util.tree_map(jp.concatenate, y_typ)
-    y_typ = jax.tree_map(lambda *x: jp.concatenate(x), *y_typ)
+    y_typ = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *y_typ)
     y_take = np.argsort(np.concatenate([key_y_take[key][i] for key in keys]))
     _check_output(y_typ, y_take, typ, i)
     y.append(_take(y_typ, y_take))

mjx/mujoco/mjx/_src/scan_test.py

@@ -55,7 +55,7 @@ class ScanTest(absltest.TestCase):
         <worldbody/>
       </mujoco>
     """)
-    m = mjx.device_put(m)
+    m = mjx.put_model(m)
 
     def fn(body_id):
       return body_id + 1
@@ -69,7 +69,7 @@ class ScanTest(absltest.TestCase):
   def test_flat_joints(self):
     """Tests scanning over bodies with joints of different types."""
     m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML)
-    m = mjx.device_put(m)
+    m = mjx.put_model(m)
 
     # we will test two functions:
     #   1) j_fn receives jnt_types as a jp array
@@ -105,7 +105,7 @@ class ScanTest(absltest.TestCase):
   def test_body_tree(self):
     """Tests tree scanning over bodies with different joint counts."""
     m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML)
-    m = mjx.device_put(m)
+    m = mjx.put_model(m)
 
     # we will test two functions:
     #   1) j_fn receives jnt_pos which is a jp array
@@ -196,7 +196,7 @@ class ScanTest(absltest.TestCase):
   def test_scan_actuators(self):
     """Tests scanning over actuators."""
     m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML)
-    m = mjx.device_put(m)
+    m = mjx.put_model(m)
 
     fn = lambda *args: args
     args = (

mjx/mujoco/mjx/_src/solver.py

@@ -17,7 +17,6 @@
 import jax
 from jax import numpy as jp
 import mujoco
-from mujoco.mjx._src import constraint
 from mujoco.mjx._src import math
 from mujoco.mjx._src import smooth
 from mujoco.mjx._src import support
@@ -81,7 +80,7 @@ class _Context(PyTreeNode):
         prev_cost=0.0,
         solver_niter=0,
     )
-    ctx = _update_constraint(m, d, ctx)
+    ctx = _update_constraint(d, ctx)
     if grad:
       ctx = _update_gradient(m, d, ctx)
       ctx = ctx.replace(search=-ctx.Mgrad)  # start with preconditioned gradient
@@ -107,7 +106,7 @@ class _LSPoint(PyTreeNode):
   @classmethod
   def create(
       cls,
-      m: Model,
+      d: Data,
       ctx: _Context,
       alpha: jax.Array,
       jv: jax.Array,
@@ -118,8 +117,7 @@ class _LSPoint(PyTreeNode):
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
 
     # TODO(robotics-team): change this to support friction constraints
-    ne, nf, *_ = constraint.count_constraints(m)
+    active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
-    active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True)
     quad = jax.vmap(jp.multiply)(quad, active)  # only active
     quad_total = quad_gauss + jp.sum(quad, axis=0)
 
@@ -161,13 +159,12 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
   return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]
 
 
-def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
+def _update_constraint(d: Data, ctx: _Context) -> _Context:
   """Updates constraint force and resulting cost given latst solver iteration.
 
   Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
 
   Args:
-    m: model defining constraints
     d: data which contains latest qacc and smooth terms
     ctx: current solver context
 
@@ -177,8 +174,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   # TODO(robotics-team): add friction constraints
 
   # only count active constraints
-  ne, nf, *_ = constraint.count_constraints(m)
+  active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
-  active = (ctx.Jaref < 0).at[:ne + nf].set(True)
 
   efc_force = d.efc_D * -ctx.Jaref * active
   qfrc_constraint = d.efc_J.T @ efc_force
@@ -217,8 +213,7 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
   if m.opt.solver == SolverType.CG:
     mgrad = smooth.solve_m(m, d, grad)
   elif m.opt.solver == SolverType.NEWTON:
-    ne, nf, *_ = constraint.count_constraints(m)
+    active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True)
-    active = (ctx.Jaref < 0).at[: ne + nf].set(True)
     h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
     h = support.full_m(m, d) + h
     h_ = jax.scipy.linalg.cho_factor(h)
@@ -262,7 +257,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
   quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
   quad = (quad * d.efc_D).T
 
-  point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss)
+  point_fn = lambda a: _LSPoint.create(d, ctx, a, jv, quad, quad_gauss)
 
   def cond(ctx: _LSContext) -> jax.Array:
     done = ctx.ls_iter >= m.opt.ls_iterations
@@ -283,14 +278,14 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
     # 1) they are not correctly at a bracket boundary (e.g. lo.deriv_0 > 0), OR
     # 2) if moving to next or mid narrows the bracket
     swap_lo_next = (lo.deriv_0 > 0) | (lo.deriv_0 < lo_next.deriv_0)
-    lo = jax.tree_map(lambda x, y: jp.where(swap_lo_next, y, x), lo, lo_next)
+    lo = jax.tree_util.tree_map(lambda x, y: jp.where(swap_lo_next, y, x), lo, lo_next)
     swap_lo_mid = (mid.deriv_0 < 0) & (lo.deriv_0 < mid.deriv_0)
-    lo = jax.tree_map(lambda x, y: jp.where(swap_lo_mid, y, x), lo, mid)
+    lo = jax.tree_util.tree_map(lambda x, y: jp.where(swap_lo_mid, y, x), lo, mid)
 
     swap_hi_next = (hi.deriv_0 < 0) | (hi.deriv_0 > hi_next.deriv_0)
-    hi = jax.tree_map(lambda x, y: jp.where(swap_hi_next, y, x), hi, hi_next)
+    hi = jax.tree_util.tree_map(lambda x, y: jp.where(swap_hi_next, y, x), hi, hi_next)
     swap_hi_mid = (mid.deriv_0 > 0) & (hi.deriv_0 > mid.deriv_0)
-    hi = jax.tree_map(lambda x, y: jp.where(swap_hi_mid, y, x), hi, mid)
+    hi = jax.tree_util.tree_map(lambda x, y: jp.where(swap_hi_mid, y, x), hi, mid)
 
     swap = swap_lo_next | swap_lo_mid | swap_hi_next | swap_hi_mid
 
@@ -302,8 +297,8 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
   p0 = point_fn(jp.array(0.0))
   lo = point_fn(p0.alpha - p0.deriv_0 / p0.deriv_1)
   lesser_fn = lambda x, y: jp.where(lo.deriv_0 < p0.deriv_0, x, y)
-  hi = jax.tree_map(lesser_fn, p0, lo)
+  hi = jax.tree_util.tree_map(lesser_fn, p0, lo)
-  lo = jax.tree_map(lesser_fn, lo, p0)
+  lo = jax.tree_util.tree_map(lesser_fn, lo, p0)
   ls_ctx = _LSContext(lo=lo, hi=hi, swap=jp.array(True), ls_iter=0)
   ls_ctx = _while_loop_scan(cond, body, ls_ctx, m.opt.ls_iterations)
 
@@ -336,7 +331,7 @@ def solve(m: Model, d: Data) -> Data:
   def body(ctx: _Context) -> _Context:
     ctx = _linesearch(m, d, ctx)
     prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad  # pylint: disable=invalid-name
-    ctx = _update_constraint(m, d, ctx)
+    ctx = _update_constraint(d, ctx)
     ctx = _update_gradient(m, d, ctx)
 
     # polak-ribiere:

mjx/mujoco/mjx/_src/support.py

@@ -208,3 +208,8 @@ def get_custom_numeric(m: Union[Model, mujoco.MjModel], name: str) -> float:
       return m.numeric_data[m.numeric_adr[i]]
 
   return -1
+
+
+def get_custom_int(m: Union[Model, mujoco.MjModel], name: str) -> int:
+  """Returns a custom integer given an MjModel or mjx.Model."""
+  return int(get_custom_numeric(m, name))

mjx/mujoco/mjx/_src/types.py

@@ -15,10 +15,8 @@
 """Base types used in MJX."""
 
 import enum
-from typing import List, Optional
 
 import jax
-import jax.numpy as jp
 import mujoco
 from mujoco.mjx._src.dataclasses import PyTreeNode  # pylint: disable=g-importing-member
 import numpy as np
@@ -213,6 +211,22 @@ class BiasType(enum.IntEnum):
   # unsupported: MUSCLE, USER
 
 
+class ConstraintType(enum.IntEnum):
+  """Type of constraint.
+
+  Attributes:
+    EQUALITY: equality constraint
+    LIMIT_JOINT: joint limit
+    CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
+  """
+  EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
+  # unsupported: FRICTION_DOF, FRICTION_TENDON
+  LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT
+  # unsupported: LIMIT_TENDON, CONTACT_FRICTIONLESS
+  CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
+  # unsupported: CONTACT_ELLIPTIC
+
+
 class CamLightType(enum.IntEnum):
   """Type of camera light.
 
@@ -371,6 +385,7 @@ class Model(PyTreeNode):
     geom_solref: constraint solver reference: contact         (ngeom, mjNREF)
     geom_solimp: constraint solver impedance: contact         (ngeom, mjNIMP)
     geom_size: geom-specific size parameters                  (ngeom, 3)
+    geom_rbound: radius of bounding sphere                    (ngeom,)
     geom_pos: local position offset rel. to body              (ngeom, 3)
     geom_quat: local orientation offset rel. to body          (ngeom, 4)
     geom_friction: friction for (slide, spin, roll)           (ngeom, 3)
@@ -388,17 +403,13 @@ class Model(PyTreeNode):
     cam_poscom0:  global position rel. to sub-com in qpos0    (ncam, 3)
     cam_pos0:  global position rel. to body in qpos0          (ncam, 3)
     cam_mat0:  global orientation in qpos0                    (ncam, 9)
-    mat_rgba: rgba                                            (nmat, 4)
+    mesh_vertadr: first vertex address                        (nmesh,)
-    mesh_vertadr: first vertex address                        (nmesh x 1)
+    mesh_faceadr: first face address                          (nmesh,)
-    mesh_faceadr: first face address                          (nmesh x 1)
+    mesh_graphadr: graph data address; -1: no graph           (nmesh,)
     mesh_vert: vertex positions for all meshes                (nmeshvert, 3)
     mesh_face: vertex face data                               (nmeshface, 3)
-    geom_convex_face: vertex face data, MJX only              (ngeom,)
+    mesh_graph: convex graph data                             (nmeshgraph,)
-    geom_convex_vert: vertex data, MJX only                   (ngeom,)
+    mat_rgba: rgba                                            (nmat, 4)
-    geom_convex_edge_dir: unique edge direction, MJX only     (ngeom,)
-    geom_convex_facenormal: normal face data, MJX only        (ngeom,)
-    geom_convex_face_edge: edges for each face                (ngeom,)
-    geom_convex_face_edge_normal: edge normals for each face  (ngeom,)
     pair_dim: contact dimensionality                          (npair,)
     pair_geom1: id of geom1                                   (npair,)
     pair_geom2: id of geom2                                   (npair,)
@@ -435,6 +446,7 @@ class Model(PyTreeNode):
     actuator_gear: scale length and transmitted force         (nu, 6)
     numeric_adr: address of field in numeric_data             (nnumeric,)
     numeric_data: array of all numeric fields                 (nnumericdata,)
+    name_meshadr: mesh name pointers                          (nmesh,)
     name_numericadr: numeric name pointers                    (nnumeric,)
     names: names of all objects, 0-terminated                 (nnames,)
   """
@@ -516,6 +528,7 @@ class Model(PyTreeNode):
   geom_solref: jax.Array
   geom_solimp: jax.Array
   geom_size: jax.Array
+  geom_rbound: jax.Array
   geom_pos: jax.Array
   geom_quat: jax.Array
   geom_friction: jax.Array
@@ -535,18 +548,14 @@ class Model(PyTreeNode):
   cam_mat0: jax.Array
   mesh_vertadr: np.ndarray
   mesh_faceadr: np.ndarray
+  mesh_graphadr: np.ndarray
   mesh_vert: np.ndarray
   mesh_face: np.ndarray
+  mesh_graph: np.ndarray
   mat_rgba: np.ndarray
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
   pair_geom2: np.ndarray
-  geom_convex_face: List[Optional[jax.Array]]
-  geom_convex_vert: List[Optional[jax.Array]]
-  geom_convex_edge_dir: List[Optional[jax.Array]]
-  geom_convex_facenormal: List[Optional[jax.Array]]
-  geom_convex_edge: List[Optional[jax.Array]]
-  geom_convex_edge_face_normal: List[Optional[jax.Array]]
   pair_solref: jax.Array
   pair_solreffriction: jax.Array
   pair_solimp: jax.Array
@@ -580,6 +589,7 @@ class Model(PyTreeNode):
   actuator_gear: jax.Array
   numeric_adr: np.ndarray
   numeric_data: np.ndarray
+  name_meshadr: np.ndarray
   name_numericadr: np.ndarray
   names: bytes
 
@@ -596,8 +606,11 @@ class Contact(PyTreeNode):
     solref: constraint solver reference, normal direction             (mjNREF,)
     solreffriction: constraint solver reference, friction directions  (mjNREF,)
     solimp: constraint solver impedance                               (mjNIMP,)
-    geom1: id of geom 1
+    dim: contact space dimensionality: 1, 3, 4, or 6
-    geom2: id of geom 2
+    geom1: id of geom 1; deprecated, use geom[0]
+    geom2: id of geom 2; deprecated, use geom[1]
+    geom: geom ids                                                    (2,)
+    efc_address: address in efc; -1: not included
   """
   dist: jax.Array
   pos: jax.Array
@@ -607,33 +620,25 @@ class Contact(PyTreeNode):
   solref: jax.Array
   solreffriction: jax.Array
   solimp: jax.Array
-  # unsupported: mu, H, dim
+  # unsupported: mu, H
+  dim: np.ndarray
   geom1: jax.Array
   geom2: jax.Array
-  # unsupported: efc_address, exclude
+  geom: jax.Array
-
+  # unsupported: flex, elem, vert, exclude
-  @classmethod
+  efc_address: np.ndarray
-  def zero(cls, ncon: int = 0) -> 'Contact':
-    """Returns a contact filled with zeros."""
-    return Contact(
-        dist=jp.zeros(ncon),
-        pos=jp.zeros((ncon, 3,)),
-        frame=jp.zeros((ncon, 3, 3)),
-        includemargin=jp.zeros(ncon),
-        friction=jp.zeros((ncon, 5)),
-        solref=jp.zeros((ncon, mujoco.mjNREF)),
-        solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
-        solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
-        geom1=jp.zeros(ncon, dtype=int),
-        geom2=jp.zeros(ncon, dtype=int),
-    )
 
 
 class Data(PyTreeNode):
-  r"""Dynamic state that updates each step.\
+  r"""Dynamic state that updates each step.
 
   Attributes:
-    solver_niter: number of solver iterations, per island         (mjNISLAND,)
+    ne: number of equality constraints
+    nf: number of friction constraints
+    nl: number of limit constraints
+    nefc: number of constraints
+    ncon: number of contacts
+    solver_niter: number of solver iterations
     time: simulation time
     qpos: position                                                (nq,)
     qvel: velocity                                                (nv,)
@@ -663,14 +668,15 @@ class Data(PyTreeNode):
     cinert: com-based body inertia and mass                       (nbody, 10)
     actuator_length: actuator lengths                             (nu,)
     actuator_moment: actuator moments                             (nu, nv)
-    crb: com-based composite inertia and mass                     (nbody, 10)
+    crb: com-based composite inertia and mass                     (nbody, 10)  \
     qM: total inertia                                  if sparse: (nM,)
                                                        if dense:  (nv, nv)
     qLD: L'*D*L (or Cholesky) factorization of M.      if sparse: (nM,)
                                                        if dense:  (nv, nv)
     qLDiagInv: 1/diag(D)                               if sparse: (nv,)
                                                        if dense:  (0,)
-    contact: list of all detected contacts                        (ncon,)
+    contact: all detected contacts                                (ncon,)
+    efc_type: constraint type                                     (nefc,)
     efc_J: constraint Jacobian                                    (nefc, nv)
     efc_frictionloss: frictionloss (friction)                     (nefc,)
     efc_D: constraint mass                                        (nefc,)
@@ -689,6 +695,12 @@ class Data(PyTreeNode):
     efc_force: constraint force in constraint space               (nefc,)
     userdata: user data, not touched by engine                    (nuserdata,)
   """
+  # constant sizes:
+  ne: int
+  nf: int
+  nl: int
+  nefc: int
+  ncon: int
   # solver statistics:
   solver_niter: jax.Array
   # global properties:
@@ -732,6 +744,7 @@ class Data(PyTreeNode):
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
   contact: Contact
+  efc_type: np.ndarray
   efc_J: jax.Array  # pylint:disable=invalid-name
   efc_frictionloss: jax.Array
   efc_D: jax.Array  # pylint:disable=invalid-name

mjx/mujoco/mjx/integration_test/collision_driver_test.py

@@ -30,7 +30,10 @@ import numpy as np
 
 def _assert_attr_eq(mjx_d, mj_d, attr, name, atol):
   if attr == 'efc_address':
-    # we do not test efc_address since it gets set in constraint logic
+    # contact order not guaranteed to match
+    np.testing.assert_array_equal(
+        np.sort(mjx_d.efc_address), np.sort(mj_d.efc_address)
+    )
     return
   err_msg = f'mismatch: {attr} in run: {name}'
   mjx_d, mj_d = getattr(mjx_d, attr), getattr(mj_d, attr)
@@ -79,9 +82,12 @@ class CollisionDriverIntegrationTest(parameterized.TestCase):
     self.assertSequenceEqual(set(idx_mjx), set(idx_mj))
     idx = sorted(range(len(idx_mj)), key=lambda x: idx_mj.index(idx_mjx[x]))
 
-    mjx_contact = jax.tree_map(
+    mjx_contact = jax.tree_util.tree_map(
         lambda x: x.take(np.array(idx), axis=0), dx.contact
     )
+    mjx_contact = mjx_contact.replace(
+        dim=mjx_contact.dim[idx], efc_address=mjx_contact.efc_address[idx]
+    )
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)
 

mjx/tutorial.ipynb

@@ -829,7 +829,7 @@
     "\n",
     "  friction, gain, bias = rand(rng)\n",
     "\n",
-    "  in_axes = jax.tree_map(lambda x: None, sys)\n",
+    "  in_axes = jax.tree_util.tree_map(lambda x: None, sys)\n",
     "  in_axes = in_axes.tree_replace({\n",
     "      'geom_friction': 0,\n",
     "      'actuator_gainprm': 0,\n",

python/mujoco/functions.cc

@@ -106,11 +106,11 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mj_defaultOption>(pymodule);
   Def<traits::mj_defaultVisual>(pymodule);
   // Skipped: mj_copyModel (have MjModel.__copy__, memory managed by MjModel)
-  DEF_WITH_OMITTED_PY_ARGS(traits::mj_saveModel, "buffer_sz")(
+  pymodule.def(
-      pymodule,
+      "mj_saveModel",
-      [](const raw::MjModel* m, const std::optional<std::string>& filename,
+      [](const MjModelWrapper& m, const std::optional<std::string>& filename = std::nullopt,
          std::optional<
-             Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>> buffer) {
+             Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>> buffer = std::nullopt) {
         void* buffer_ptr = nullptr;
         int buffer_sz = 0;
         if (buffer.has_value()) {
@@ -118,9 +118,13 @@ PYBIND11_MODULE(_functions, pymodule) {
           buffer_sz = buffer->size();
         }
         return InterceptMjErrors(::mj_saveModel)(
-            m, filename.has_value() ? filename->c_str() : nullptr,
+            m.get(), filename.has_value() ? filename->c_str() : nullptr,
             buffer_ptr, buffer_sz);
-      });
+      },
+      py::arg("m"), py::arg_v("filename", std::nullopt),
+      py::arg_v("buffer", std::nullopt),
+      py::doc(traits::mj_saveModel::doc),
+      py::call_guard<py::gil_scoped_release>());
   // Skipped: mj_loadModel (have MjModel.from_binary_path)
   // Skipped: mj_deleteModel (have MjModel.__del__)
   Def<traits::mj_sizeModel>(pymodule);

python/mujoco/renderer.py

@@ -78,8 +78,11 @@ the clause:
 
     # Create render contexts.
     # TODO(nimrod): Figure out why pytype doesn't like gl_context.GLContext
-    self._gl_context = gl_context.GLContext(width, height)  # type: ignore
+    self._gl_context = None  # type: ignore
-    self._gl_context.make_current()
+    if gl_context.GLContext is not None:
+      self._gl_context = gl_context.GLContext(width, height)
+    if self._gl_context:
+      self._gl_context.make_current()
     self._mjr_context = _render.MjrContext(
         model, _enums.mjtFontScale.mjFONTSCALE_150.value
     )
@@ -148,9 +151,11 @@ the clause:
       self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True
       self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True
 
-    if self._gl_context is None:
+    if self._mjr_context is None:
       raise RuntimeError('render cannot be called after close.')
-    self._gl_context.make_current()
+
+    if self._gl_context:
+      self._gl_context.make_current()
 
     if self._depth_rendering:
       out_shape = (self._height, self._width)

src/user/user_api.cc

@@ -91,6 +91,17 @@ int mjs_attachBody(mjsFrame* parent, const mjsBody* child,
 
 
 
+// attach frame to a parent body
+int mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
+                    const char* prefix, const char* suffix) {
+  mjCBody* body_parent = static_cast<mjCBody*>(parent->element);
+  mjCFrame* child_frame = static_cast<mjCFrame*>(child->element);
+  *body_parent += std::string(prefix) + *child_frame + std::string(suffix);
+  return 0;
+}
+
+
+
 // get error message from model
 const char* mjs_getError(mjSpec* s) {
   mjCModel* modelC = static_cast<mjCModel*>(s->element);

src/user/user_api.h

@@ -40,39 +40,39 @@ typedef struct _mjDoubleVec* mjDoubleVec;
 
 //---------------------------------- enum types (mjt) ----------------------------------------------
 
-typedef enum _mjtGeomInertia {    // type of inertia inference
+typedef enum _mjtGeomInertia {     // type of inertia inference
-  mjINERTIA_VOLUME,               // mass distributed in the volume
+  mjINERTIA_VOLUME,                // mass distributed in the volume
-  mjINERTIA_SHELL,                // mass distributed on the surface
+  mjINERTIA_SHELL,                 // mass distributed on the surface
 } mjtGeomInertia;
 
 
-typedef enum _mjtBuiltin {        // type of built-in procedural texture
+typedef enum _mjtBuiltin {         // type of built-in procedural texture
-  mjBUILTIN_NONE = 0,             // no built-in texture
+  mjBUILTIN_NONE = 0,              // no built-in texture
-  mjBUILTIN_GRADIENT,             // gradient: rgb1->rgb2
+  mjBUILTIN_GRADIENT,              // gradient: rgb1->rgb2
-  mjBUILTIN_CHECKER,              // checker pattern: rgb1, rgb2
+  mjBUILTIN_CHECKER,               // checker pattern: rgb1, rgb2
-  mjBUILTIN_FLAT                  // 2d: rgb1; cube: rgb1-up, rgb2-side, rgb3-down
+  mjBUILTIN_FLAT                   // 2d: rgb1; cube: rgb1-up, rgb2-side, rgb3-down
 } mjtBuiltin;
 
 
-typedef enum _mjtMark {           // mark type for procedural textures
+typedef enum _mjtMark {            // mark type for procedural textures
-  mjMARK_NONE = 0,                // no mark
+  mjMARK_NONE = 0,                 // no mark
-  mjMARK_EDGE,                    // edges
+  mjMARK_EDGE,                     // edges
-  mjMARK_CROSS,                   // cross
+  mjMARK_CROSS,                    // cross
-  mjMARK_RANDOM                   // random dots
+  mjMARK_RANDOM                    // random dots
 } mjtMark;
 
 
-typedef enum _mjtLimited {        // type of limit specification
+typedef enum _mjtLimited {         // type of limit specification
-  mjLIMITED_FALSE = 0,            // not limited
+  mjLIMITED_FALSE = 0,             // not limited
-  mjLIMITED_TRUE,                 // limited
+  mjLIMITED_TRUE,                  // limited
-  mjLIMITED_AUTO,                 // limited inferred from presence of range
+  mjLIMITED_AUTO,                  // limited inferred from presence of range
 } mjtLimited;
 
 
 typedef enum _mjtInertiaFromGeom {
-  mjINERTIAFROMGEOM_FALSE = 0,    // do not use; inertial element required
+  mjINERTIAFROMGEOM_FALSE = 0,     // do not use; inertial element required
-  mjINERTIAFROMGEOM_TRUE,         // always use; overwrite inertial element
+  mjINERTIAFROMGEOM_TRUE,          // always use; overwrite inertial element
-  mjINERTIAFROMGEOM_AUTO          // use only if inertial element is missing
+  mjINERTIAFROMGEOM_AUTO           // use only if inertial element is missing
 } mjtInertiaFromGeom;
 
 
@@ -84,8 +84,8 @@ typedef struct _mjElement {        // element type, do not modify
 
 
 typedef struct _mjSpec {           // model specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
-  mjStatistic stat;                // statistics override (if defined)
+  mjString modelname;              // model name
 
   // compiler settings
   mjtByte autolimits;              // infer "limited" attribute based on range
@@ -109,33 +109,33 @@ typedef struct _mjSpec {           // model specification
   mjLROpt LRopt;                   // options for lengthrange computation
 
   // engine data
-  mjString modelname;             // model name
+  mjOption option;                 // physics options
-  mjOption option;                // options
+  mjVisual visual;                 // visual options
-  mjVisual visual;                // visual options
+  mjStatistic stat;                // statistics override (if defined)
-  size_t memory;                  // size of arena+stack memory in bytes
-  int nemax;                      // max number of equality constraints
-  int njmax;                      // max number of constraints (Jacobian rows)
-  int nconmax;                    // max number of detected contacts (mjContact array size)
-  size_t nstack;                  // (deprecated) number of fields in mjData stack
-  int nuserdata;                  // number extra fields in mjData
-  int nuser_body;                 // number of mjtNums in body_user
-  int nuser_jnt;                  // number of mjtNums in jnt_user
-  int nuser_geom;                 // number of mjtNums in geom_user
-  int nuser_site;                 // number of mjtNums in site_user
-  int nuser_cam;                  // number of mjtNums in cam_user
-  int nuser_tendon;               // number of mjtNums in tendon_user
-  int nuser_actuator;             // number of mjtNums in actuator_user
-  int nuser_sensor;               // number of mjtNums in sensor_user
 
   // sizes
-  int nkey;                       // number of keyframes
+  size_t memory;                   // number of bytes in arena+stack memory
+  int nemax;                       // max number of equality constraints
+  int nuserdata;                   // number of mjtNums in userdata
+  int nuser_body;                  // number of mjtNums in body_user
+  int nuser_jnt;                   // number of mjtNums in jnt_user
+  int nuser_geom;                  // number of mjtNums in geom_user
+  int nuser_site;                  // number of mjtNums in site_user
+  int nuser_cam;                   // number of mjtNums in cam_user
+  int nuser_tendon;                // number of mjtNums in tendon_user
+  int nuser_actuator;              // number of mjtNums in actuator_user
+  int nuser_sensor;                // number of mjtNums in sensor_user
+  int nkey;                        // number of keyframes
+  int njmax;                       // (deprecated) max number of constraints
+  int nconmax;                     // (deprecated) max number of detected contacts
+  size_t nstack;                   // (deprecated) number of mjtNums in mjData stack
 
   // global data
-  mjString comment;               // comment at top of XML
+  mjString comment;                // comment at top of XML
-  mjString modelfiledir;          // path to model file
+  mjString modelfiledir;           // path to model file
 
   // other
-  bool hasImplicitPluginElem;     // already encountered an implicit plugin sensor/actuator
+  bool hasImplicitPluginElem;      // already encountered an implicit plugin sensor/actuator
 } mjSpec;
 
 
@@ -148,7 +148,7 @@ typedef struct _mjsOrientation {   // alternative orientation specifiers
 
 
 typedef struct _mjsPlugin {        // plugin specification
-  mjElement* instance;             // object type
+  mjElement* instance;             // element type
   mjString name;                   // name
   mjString instance_name;          // instance name
   int plugin_slot;                 // global registered slot number of the plugin
@@ -158,7 +158,7 @@ typedef struct _mjsPlugin {        // plugin specification
 
 
 typedef struct _mjsBody {          // body specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString childclass;             // childclass name
 
@@ -186,7 +186,7 @@ typedef struct _mjsBody {          // body specification
 
 
 typedef struct _mjsFrame {         // frame specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString childclass;             // childclass name
   double pos[3];                   // position
@@ -197,7 +197,7 @@ typedef struct _mjsFrame {         // frame specification
 
 
 typedef struct _mjsJoint {         // joint specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjtJoint type;                   // joint type
@@ -238,7 +238,7 @@ typedef struct _mjsJoint {         // joint specification
 
 
 typedef struct _mjsGeom {          // geom specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // classname
   mjtGeom type;                    // geom type
@@ -287,7 +287,7 @@ typedef struct _mjsGeom {          // geom specification
 
 
 typedef struct _mjsSite {          // site specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -311,7 +311,7 @@ typedef struct _mjsSite {          // site specification
 
 
 typedef struct _mjsCamera {        // camera specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -340,7 +340,7 @@ typedef struct _mjsCamera {        // camera specification
 
 
 typedef struct _mjsLight {         // light specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -368,7 +368,7 @@ typedef struct _mjsLight {         // light specification
 
 
 typedef struct _mjsFlex {
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -409,7 +409,7 @@ typedef struct _mjsFlex {
 
 
 typedef struct _mjsMesh {          // mesh specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjString content_type;           // content type of file
@@ -430,7 +430,7 @@ typedef struct _mjsMesh {          // mesh specification
 
 
 typedef struct _mjsHField {        // height field specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString content_type;           // content type of file
   mjString file;                   // file: (nrow, ncol, [elevation data])
@@ -444,7 +444,7 @@ typedef struct _mjsHField {        // height field specification
 
 
 typedef struct _mjsSkin {          // skin specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjString file;                   // skin file
@@ -471,7 +471,7 @@ typedef struct _mjsSkin {          // skin specification
 
 
 typedef struct _mjsTexture {       // texture specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjtTexture type;                 // texture type
@@ -505,7 +505,7 @@ typedef struct _mjsTexture {       // texture specification
 
 
 typedef struct _mjsMaterial {      // material specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjString texture;                // name of texture (empty: none)
@@ -521,7 +521,7 @@ typedef struct _mjsMaterial {      // material specification
 
 
 typedef struct _mjsPair {
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjString geomname1;              // name of geom 1
@@ -540,7 +540,7 @@ typedef struct _mjsPair {
 
 
 typedef struct _mjsExclude {
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString bodyname1;              // name of geom 1
   mjString bodyname2;              // name of geom 2
@@ -549,7 +549,7 @@ typedef struct _mjsExclude {
 
 
 typedef struct _mjsEquality {      // equality specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
   mjtEq type;                      // constraint type
@@ -564,7 +564,7 @@ typedef struct _mjsEquality {      // equality specification
 
 
 typedef struct _mjsTendon {        // tendon specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -596,13 +596,13 @@ typedef struct _mjsTendon {        // tendon specification
 
 
 typedef struct _mjsWrap {          // wrapping object specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString info;                   // message appended to errors
 } mjsWrap;
 
 
 typedef struct _mjsActuator {      // actuator specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -646,7 +646,7 @@ typedef struct _mjsActuator {      // actuator specification
 
 
 typedef struct _mjsSensor {        // sensor specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString classname;              // class name
 
@@ -674,7 +674,7 @@ typedef struct _mjsSensor {        // sensor specification
 
 
 typedef struct _mjsNumeric {       // custom numeric field specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjDoubleVec data;                // initialization data
   int size;                        // array size, can be bigger than data size
@@ -683,7 +683,7 @@ typedef struct _mjsNumeric {       // custom numeric field specification
 
 
 typedef struct _mjsText {          // custom text specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjString data;                   // text string
   mjString info;                   // message appended to compiler errors
@@ -691,7 +691,7 @@ typedef struct _mjsText {          // custom text specification
 
 
 typedef struct _mjsTuple {         // tuple specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   mjIntVec objtype;                // object types
   mjStringVec objname;             // object names
@@ -701,7 +701,7 @@ typedef struct _mjsTuple {         // tuple specification
 
 
 typedef struct _mjsKey {           // keyframe specification
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjString name;                   // name
   double time;                     // time
   mjDoubleVec qpos;                // qpos
@@ -716,7 +716,7 @@ typedef struct _mjsKey {           // keyframe specification
 
 typedef struct _mjsDefault {       // default specification
   mjString name;                   // name
-  mjElement* element;              // object type
+  mjElement* element;              // element type
   mjsJoint* joint;                 // joint defaults
   mjsGeom* geom;                   // geom defaults
   mjsSite* site;                   // site defaults
@@ -732,34 +732,44 @@ typedef struct _mjsDefault {       // default specification
 } mjsDefault;
 
 
-//---------------------------------- API functions -------------------------------------------------
+//---------------------------------- Top-level spec manipulation -----------------------------------
 
-// Create model.
+// Create spec.
 MJAPI mjSpec* mjs_createSpec();
 
-// Copy model.
+// Compile spec to model.
-MJAPI mjSpec* mjs_copySpec(const mjSpec* s);
-
-// Copy back model.
-MJAPI void mjs_copyBack(mjSpec* s, const mjModel* m);
-
-// Compile model.
 MJAPI mjModel* mjs_compile(mjSpec* s, const mjVFS* vfs);
 
-// Attach child body to a frame of the parent, return 0 if success
+// Copy spec.
+MJAPI mjSpec* mjs_copySpec(const mjSpec* s);
+
+// Get compiler error message from spec.
+MJAPI const char* mjs_getError(mjSpec* s);
+
+// Return 1 if compiler error is a warning.
+MJAPI int mjs_isWarning(mjSpec* s);
+
+// Copy model fields back into spec.
+MJAPI void mjs_copyBack(mjSpec* s, const mjModel* m);
+
+// Delete spec.
+MJAPI void mjs_deleteSpec(mjSpec* s);
+
+
+//---------------------------------- Attachment ----------------------------------------------------
+
+// Attach child body to a parent frame, return 0 on success.
 MJAPI int mjs_attachBody(mjsFrame* parent, const mjsBody* child,
                          const char* prefix, const char* suffix);
 
-// Get error message from model.
+// Attach child frame to a parent body, return 0 if success.
-MJAPI const char* mjs_getError(mjSpec* s);
+MJAPI int mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
+                          const char* prefix, const char* suffix);
 
-// Return 1 if model has warnings.
-MJAPI int mjs_isWarning(mjSpec* s);
 
-// Delete model.
+//---------------------------------- Add tree elements ---------------------------------------------
-MJAPI void mjs_deleteSpec(mjSpec* s);
 
-// Add child body to body, return child spec.
+// Add child body to body, return child.
 MJAPI mjsBody* mjs_addBody(mjsBody* body, mjsDefault* def);
 
 // Add site to body, return site spec.
@@ -783,34 +793,28 @@ MJAPI mjsLight* mjs_addLight(mjsBody* body, mjsDefault* def);
 // Add frame to body.
 MJAPI mjsFrame* mjs_addFrame(mjsBody* body, mjsFrame* parentframe);
 
-// Add flex to model.
+
+//---------------------------------- Add non-tree elements -----------------------------------------
+
+// Add actuator.
+MJAPI mjsActuator* mjs_addActuator(mjSpec* s, mjsDefault* def);
+
+// Add sensor.
+MJAPI mjsSensor* mjs_addSensor(mjSpec* s);
+
+// Add flex.
 MJAPI mjsFlex* mjs_addFlex(mjSpec* s);
 
-// Add mesh to model.
+// Add contact pair.
-MJAPI mjsMesh* mjs_addMesh(mjSpec* s, mjsDefault* def);
-
-// Add height field to model.
-MJAPI mjsHField* mjs_addHField(mjSpec* s);
-
-// Add skin to model.
-MJAPI mjsSkin* mjs_addSkin(mjSpec* s);
-
-// Add texture to model.
-MJAPI mjsTexture* mjs_addTexture(mjSpec* s);
-
-// Add material to model.
-MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
-
-// Add pair to model.
 MJAPI mjsPair* mjs_addPair(mjSpec* s, mjsDefault* def);
 
-// Add excluded body pair to model.
+// Add excluded body pair.
 MJAPI mjsExclude* mjs_addExclude(mjSpec* s);
 
-// Add equality to model.
+// Add equality.
 MJAPI mjsEquality* mjs_addEquality(mjSpec* s, mjsDefault* def);
 
-// Add tendon to model.
+// Add tendon.
 MJAPI mjsTendon* mjs_addTendon(mjSpec* s, mjsDefault* def);
 
 // Wrap site using tendon.
@@ -825,42 +829,48 @@ MJAPI mjsWrap* mjs_wrapJoint(mjsTendon* tendon, const char* name, double coef);
 // Wrap pulley using tendon.
 MJAPI mjsWrap* mjs_wrapPulley(mjsTendon* tendon, double divisor);
 
-// Add actuator to model.
+// Add numeric.
-MJAPI mjsActuator* mjs_addActuator(mjSpec* s, mjsDefault* def);
-
-// Add sensor to model.
-MJAPI mjsSensor* mjs_addSensor(mjSpec* s);
-
-// Add numeric to model.
 MJAPI mjsNumeric* mjs_addNumeric(mjSpec* s);
 
-// Add text to model.
+// Add text.
 MJAPI mjsText* mjs_addText(mjSpec* s);
 
-// Add tuple to model.
+// Add tuple.
 MJAPI mjsTuple* mjs_addTuple(mjSpec* s);
 
-// Add keyframe to model.
+// Add keyframe.
 MJAPI mjsKey* mjs_addKey(mjSpec* s);
 
-// Add plugin to model.
+// Add plugin.
 MJAPI mjsPlugin* mjs_addPlugin(mjSpec* s);
 
-// Add default to model.
+// Add default.
 MJAPI mjsDefault* mjs_addDefault(mjSpec* s, const char* classname, int parentid, int* id);
 
-// Get model spec from body.
+
+//---------------------------------- Add assets ----------------------------------------------------
+
+// Add mesh.
+MJAPI mjsMesh* mjs_addMesh(mjSpec* s, mjsDefault* def);
+
+// Add height field.
+MJAPI mjsHField* mjs_addHField(mjSpec* s);
+
+// Add skin.
+MJAPI mjsSkin* mjs_addSkin(mjSpec* s);
+
+// Add texture.
+MJAPI mjsTexture* mjs_addTexture(mjSpec* s);
+
+// Add material.
+MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
+
+
+//---------------------------------- Find/get utilities --------------------------------------------
+
+// Get spec from body.
 MJAPI mjSpec* mjs_getSpec(mjsBody* body);
 
-// Get default corresponding to an mjElement.
-MJAPI mjsDefault* mjs_getDefault(mjElement* element);
-
-// Find default in model by class name.
-MJAPI mjsDefault* mjs_findDefault(mjSpec* s, const char* classname);
-
-// Get global default from model.
-MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s);
-
 // Find body in model by name.
 MJAPI mjsBody* mjs_findBody(mjSpec* s, const char* name);
 
@@ -873,9 +883,21 @@ MJAPI mjsMesh* mjs_findMesh(mjSpec* s, const char* name);
 // Find frame by name.
 MJAPI mjsFrame* mjs_findFrame(mjSpec* s, const char* name);
 
+// Get default corresponding to an element.
+MJAPI mjsDefault* mjs_getDefault(mjElement* element);
+
+// Find default in model by class name.
+MJAPI mjsDefault* mjs_findDefault(mjSpec* s, const char* classname);
+
+// Get global default from model.
+MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s);
+
 // Get element id.
 MJAPI int mjs_getId(mjElement* element);
 
+
+//---------------------------------- Attribute setters ---------------------------------------------
+
 // Copy text to string.
 MJAPI void mjs_setString(mjString dest, const char* text);
 
@@ -903,33 +925,39 @@ MJAPI void mjs_appendFloatVec(mjFloatVecVec dest, const float* array, int size);
 // Copy double array to vector.
 MJAPI void mjs_setDouble(mjDoubleVec dest, const double* array, int size);
 
+// Set plugin attributes.
+MJAPI void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes);
+
+
+//---------------------------------- Attribute getters ---------------------------------------------
+
 // Get string contents.
 MJAPI const char* mjs_getString(mjString source);
 
 // Get double array contents and optionally its size.
 MJAPI const double* mjs_getDouble(mjDoubleVec source, int* size);
 
-// Set plugin attributes.
+
-MJAPI void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes);
+//---------------------------------- Other utilities -----------------------------------------------
 
 // Set active plugins.
 MJAPI void mjs_setActivePlugins(mjSpec* s, void* activeplugins);
 
-// Set default.
+// Set element's default.
 MJAPI void mjs_setDefault(mjElement* element, mjsDefault* def);
 
-// Set frame.
+// Set element's enlcosing frame.
 MJAPI void mjs_setFrame(mjElement* dest, mjsFrame* frame);
 
 // Resolve alternative orientations to quat.
 MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
                                          const mjsOrientation* orientation);
 
-// Compute quat and inertia from body->fullinertia.
+// Compute quat and inertia from fullinertia, return error if any.
 MJAPI const char* mjs_setFullInertia(mjsBody* body, double quat[4], double inertia[3]);
 
 
-//---------------------------------- Initialization functions --------------------------------------
+//---------------------------------- Initialization  -----------------------------------------------
 
 // Default model attributes.
 MJAPI void mjs_defaultSpec(mjSpec& model);
@@ -1006,7 +1034,8 @@ MJAPI void mjs_defaultKey(mjsKey& key);
 // Default plugin attributes.
 MJAPI void mjs_defaultPlugin(mjsPlugin& plugin);
 
-//------------------------- Cache functions ------------------------------------
+
+//---------------------------------- Compiler cache ------------------------------------------------
 
 typedef struct _mjCache* mjCache;
 

src/user/user_model.cc

@@ -3118,7 +3118,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
 
   // compile objects in kinematic tree
   for (int i=0; i<bodies.size(); i++) {
-    bodies[i]->Compile();  // also compiles joints, geoms, sites, cameras, lights
+    bodies[i]->Compile();  // also compiles joints, geoms, sites, cameras, lights, frames
   }
 
   // compile all other objects except for keyframes

src/user/user_objects.cc

@@ -628,7 +628,6 @@ void mjCBase::SetFrame(mjCFrame* _frame) {
     return;
   }
   frame = _frame;
-  frame->Compile();
 }
 
 
@@ -729,13 +728,67 @@ mjCBody& mjCBody::operator+=(const mjCBody& other) {
 
 
 
+// attach frame to body
+mjCBody& mjCBody::operator+=(const mjCFrame& other) {
+  mjCBody* subtree = other.body;
+  other.model->prefix = other.prefix;
+  other.model->suffix = other.suffix;
+
+  // copy input frame
+  frames.push_back(new mjCFrame(other));
+  frames.back()->body = this;
+  frames.back()->model = model;
+  frames.back()->frame = other.frame;
+  int i = frames.size();
+
+  // map input frames to index in this->frames
+  std::map<mjCFrame*, int> fmap;
+  for (auto frame : subtree->frames) {
+    if (frame == static_cast<const mjCFrame*>(&other)) {
+      fmap[frame] = frames.size() - 1;
+    } else if (other.IsAncestor(frame)) {
+      fmap[frame] = i++;
+    }
+  }
+
+  // copy children that are inside the input frame
+  CopyList(frames, subtree->frames, fmap, &other);  // needs to be done first
+  CopyList(geoms, subtree->geoms, fmap, &other);
+  CopyList(joints, subtree->joints, fmap, &other);
+  CopyList(sites, subtree->sites, fmap, &other);
+  CopyList(cameras, subtree->cameras, fmap, &other);
+  CopyList(lights, subtree->lights, fmap, &other);
+
+  for (int i=0; i<subtree->bodies.size(); i++) {
+    if (!other.IsAncestor(subtree->bodies[i]->frame)) {
+      continue;
+    }
+    bodies.push_back(new mjCBody(*subtree->bodies[i], model));  // triggers recursive call
+    bodies.back()->frame =
+        subtree->bodies[i]->frame ? frames[fmap[subtree->bodies[i]->frame]] : nullptr;
+  }
+
+  // name space
+  this->NameSpace(other.model);
+
+  // attach referencing elements
+  *model += *other.model;
+
+  // clear namespace and return body
+  other.model->prefix.clear();
+  other.model->suffix.clear();
+  return *this;
+}
+
+
+
 // copy src list of elements into dst; set body, model and frame
 template <typename T>
 void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
                        std::map<mjCFrame*, int>& fmap, const mjCFrame* pframe) {
   int nsrc = (int)src.size();
   for (int i=0; i<nsrc; i++) {
-    if (pframe && src[i]->frame != pframe) {
+    if (pframe && !pframe->IsAncestor(src[i]->frame)) {
       continue;  // skip if the element is not inside pframe
     }
     dst.push_back(new T(*src[i]));
@@ -1169,6 +1222,11 @@ void mjCBody::MakeInertialExplicit() {
 void mjCBody::Compile(void) {
   CopyFromSpec();
 
+  // compile all frames
+  for (int i=0; i<frames.size(); i++) {
+    frames[i]->Compile();
+  }
+
   // resize userdata
   if (userdata_.size() > model->nuser_body) {
     throw mjCError(this, "user has more values than nuser_body in body '%s' (id = %d)",
@@ -1417,6 +1475,21 @@ mjCFrame& mjCFrame::operator+=(const mjCBody& other) {
 
 
 
+// return true if child is descendent of this frame
+bool mjCFrame::IsAncestor(const mjCFrame* child) const {
+  if (!child) {
+    return false;
+  }
+
+  if (child == this) {
+    return true;
+  }
+
+  return IsAncestor(child->frame);
+}
+
+
+
 void mjCFrame::SetParent(mjCBody* _body) {
   body = _body;
 }

src/user/user_objects.h

@@ -275,6 +275,7 @@ class mjCBody : public mjCBody_, private mjsBody {
 
   // API for adding existing objects to body
   mjCBody& operator+=(const mjCBody& other);
+  mjCBody& operator+=(const mjCFrame& other);
 
   // API for accessing objects
   int NumObjects(mjtObj type);
@@ -360,7 +361,8 @@ class mjCFrame : public mjCFrame_, private mjsFrame {
   void SetParent(mjCBody* _body);
 
   mjCFrame& operator+=(const mjCBody& other);
-  mjCFrame& operator+=(const mjCFrame& other);
+
+  bool IsAncestor(const mjCFrame* child) const;  // true if child is contained in this frame
 
  private:
   mjCFrame(mjCModel* = 0, mjCFrame* = 0);      // constructor

test/user/user_api_test.cc

@@ -196,20 +196,20 @@ TEST_F(PluginTest, RecompileCompareCache) {
 }
 
 // -------------------------------- test attach -------------------------------
-TEST_F(MujocoTest, AttachSame) {
+static constexpr char xml_child[] = R"(
-  std::array<char, 1000> er;
-  mjtNum tol = 0;
-  std::string field = "";
-
-  static constexpr char xml[] = R"(
   <mujoco>
     <worldbody>
-      <body name="body">
+      <frame name="pframe">
-        <joint type="hinge" name="hinge"/>
+        <frame name="cframe">
-        <geom type="cylinder" size=".1 1 0"/>
+          <body name="body">
-        <light mode="targetbody" target="child"/>
+            <joint type="hinge" name="hinge"/>
-        <body name="child"/>
+            <geom type="cylinder" size=".1 1 0"/>
-      </body>
+            <light mode="targetbody" target="targetbody"/>
+            <body name="targetbody"/>
+            <body/>
+          </body>
+        </frame>
+      </frame>
       <body name="ignore"/>
       <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
     </worldbody>
@@ -227,26 +227,33 @@ TEST_F(MujocoTest, AttachSame) {
       <position name="fixed" tendon="fixed"/>
     </actuator>
     <contact>
-      <exclude body1="body" body2="child"/>
+      <exclude body1="body" body2="targetbody"/>
     </contact>
   </mujoco>)";
 
+TEST_F(MujocoTest, AttachSame) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
   static constexpr char xml_result[] = R"(
   <mujoco>
     <worldbody>
       <body name="body">
         <joint type="hinge" name="hinge"/>
         <geom type="cylinder" size=".1 1 0"/>
-        <light mode="targetbody" target="child"/>
+        <light mode="targetbody" target="targetbody"/>
-        <body name="child"/>
+        <body name="targetbody"/>
+        <body/>
       </body>
       <body name="ignore"/>
       <frame name="frame" pos=".1 0 0" euler="0 90 0">
         <body name="attached-body-1">
           <joint type="hinge" name="attached-hinge-1"/>
           <geom type="cylinder" size=".1 1 0"/>
-          <light mode="targetbody" target="attached-child-1"/>
+          <light mode="targetbody" target="attached-targetbody-1"/>
-          <body name="attached-child-1"/>
+          <body name="attached-targetbody-1"/>
+          <body/>
         </body>
       </frame>
     </worldbody>
@@ -270,13 +277,13 @@ TEST_F(MujocoTest, AttachSame) {
       <position name="attached-fixed-1" tendon="attached-fixed-1"/>
     </actuator>
     <contact>
-      <exclude body1="body" body2="child"/>
+      <exclude body1="body" body2="targetbody"/>
-      <exclude body1="attached-body-1" body2="attached-child-1"/>
+      <exclude body1="attached-body-1" body2="attached-targetbody-1"/>
     </contact>
   </mujoco>)";
 
   // create parent
-  mjSpec* parent = ParseSpecFromString(xml, er.data(), er.size());
+  mjSpec* parent = ParseSpecFromString(xml_child, er.data(), er.size());
   EXPECT_THAT(parent, NotNull()) << er.data();
 
   // get frame
@@ -296,10 +303,10 @@ TEST_F(MujocoTest, AttachSame) {
   EXPECT_THAT(m_attached, NotNull());
 
   // check full name stored in mjModel
-  EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 4), "attached-body-1");
+  EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 5), "attached-body-1");
 
   // check body 3 is attached to the world
-  EXPECT_THAT(m_attached->body_parentid[3], 0);
+  EXPECT_THAT(m_attached->body_parentid[4], 0);
 
   // compare with expected XML
   mjModel* m_expected = LoadModelFromString(xml_result, er.data(), er.size());
@@ -330,32 +337,6 @@ TEST_F(MujocoTest, AttachDifferent) {
     </worldbody>
   </mujoco>)";
 
-  static constexpr char xml_child[] = R"(
-  <mujoco>
-    <worldbody>
-      <body name="body">
-        <joint type="hinge" name="hinge"/>
-        <geom type="cylinder" size=".1 1 0"/>
-          <body name="named"/>
-          <body/>
-      </body>
-      <body name="discard"/>
-    </worldbody>
-    <sensor>
-      <framepos name="keep" objtype="body" objname="body"/>
-      <framepos name="discard" objtype="body" objname="discard"/>
-    </sensor>
-    <tendon>
-      <fixed name="fixed">
-        <joint joint="hinge" coef="2"/>
-      </fixed>
-    </tendon>
-    <actuator>
-      <position joint="hinge"/>
-      <position tendon="fixed"/>
-    </actuator>
-  </mujoco>)";
-
   static constexpr char xml_result[] = R"(
   <mujoco>
     <worldbody>
@@ -366,14 +347,15 @@ TEST_F(MujocoTest, AttachDifferent) {
           <body name="attached-body-1">
             <joint type="hinge" name="attached-hinge-1"/>
             <geom type="cylinder" size=".1 1 0"/>
-              <body name="attached-named-1"/>
+            <light mode="targetbody" target="attached-targetbody-1"/>
+              <body name="attached-targetbody-1"/>
               <body/>
           </body>
         </frame>
       </body>
     </worldbody>
     <sensor>
-      <framepos name="attached-keep-1" objtype="body" objname="attached-body-1"/>
+      <framepos name="attached-sensor-1" objtype="body" objname="attached-body-1"/>
     </sensor>
     <tendon>
       <fixed name="attached-fixed-1">
@@ -381,9 +363,12 @@ TEST_F(MujocoTest, AttachDifferent) {
       </fixed>
     </tendon>
     <actuator>
-      <position joint="attached-hinge-1"/>
+      <position name="attached-hinge-1" joint="attached-hinge-1"/>
-      <position tendon="attached-fixed-1"/>
+      <position name="attached-fixed-1" tendon="attached-fixed-1"/>
     </actuator>
+    <contact>
+      <exclude body1="attached-body-1" body2="attached-targetbody-1"/>
+    </contact>
   </mujoco>)";
 
   // model with one free sphere and a frame
@@ -403,7 +388,7 @@ TEST_F(MujocoTest, AttachDifferent) {
   EXPECT_THAT(body, NotNull());
 
   // attach child to parent frame
-  EXPECT_THAT(
+  EXPECT_EQ(
       mjs_attachBody(frame, body, /*prefix=*/"attached-", /*suffix=*/"-1"), 0);
 
   // compile new model
@@ -430,5 +415,102 @@ TEST_F(MujocoTest, AttachDifferent) {
   mj_deleteModel(m_expected);
 }
 
+TEST_F(MujocoTest, AttachFrame) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
+  static constexpr char xml_parent[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="sphere">
+        <freejoint/>
+        <geom size=".1"/>
+        <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
+      </body>
+    </worldbody>
+  </mujoco>)";
+
+  static constexpr char xml_result[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="sphere">
+        <freejoint/>
+        <geom size=".1"/>
+        <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
+        <frame name="pframe">
+          <frame name="cframe">
+            <body name="attached-body-1">
+              <joint type="hinge" name="attached-hinge-1"/>
+              <geom type="cylinder" size=".1 1 0"/>
+              <light mode="targetbody" target="attached-targetbody-1"/>
+                <body name="attached-targetbody-1"/>
+                <body/>
+            </body>
+          </frame>
+        </frame>
+      </body>
+    </worldbody>
+    <sensor>
+      <framepos name="attached-sensor-1" objtype="body" objname="attached-body-1"/>
+    </sensor>
+    <tendon>
+      <fixed name="attached-fixed-1">
+        <joint joint="attached-hinge-1" coef="2"/>
+      </fixed>
+    </tendon>
+    <actuator>
+      <position name="attached-hinge-1" joint="attached-hinge-1"/>
+      <position name="attached-fixed-1" tendon="attached-fixed-1"/>
+    </actuator>
+    <contact>
+      <exclude body1="attached-body-1" body2="attached-targetbody-1"/>
+    </contact>
+  </mujoco>)";
+
+  // model with one free sphere and a frame
+  mjSpec* parent = ParseSpecFromString(xml_parent, er.data(), er.size());
+  EXPECT_THAT(parent, NotNull()) << er.data();
+
+  // get frame
+  mjsBody* body = mjs_findBody(parent, "sphere");
+  EXPECT_THAT(body, NotNull());
+
+  // model with one cylinder and a hinge
+  mjSpec* child = ParseSpecFromString(xml_child, er.data(), er.size());
+  EXPECT_THAT(child, NotNull()) << er.data();
+
+  // get subtree
+  mjsFrame* frame = mjs_findFrame(child, "pframe");
+  EXPECT_THAT(frame, NotNull());
+
+  // attach child to parent frame
+  EXPECT_THAT(
+      mjs_attachFrame(body, frame, /*prefix=*/"attached-", /*suffix=*/"-1"), 0);
+
+  // compile new model
+  mjModel* m_attached = mjs_compile(parent, 0);
+  EXPECT_THAT(m_attached, NotNull());
+
+  // check full name stored in mjModel
+  EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 2), "attached-body-1");
+
+  // check body 2 is attached to body 1
+  EXPECT_THAT(m_attached->body_parentid[2], 1);
+
+  // compare with expected XML
+  mjModel* m_expected = LoadModelFromString(xml_result, er.data(), er.size());
+  EXPECT_THAT(m_expected, NotNull()) << er.data();
+  EXPECT_LE(CompareModel(m_attached, m_expected, field), tol)
+            << "Expected and attached models are different!\n"
+            << "Different field: " << field << '\n';;
+
+  // destroy everything
+  mjs_deleteSpec(parent);
+  mjs_deleteSpec(child);
+  mj_deleteModel(m_attached);
+  mj_deleteModel(m_expected);
+}
+
 }  // namespace
 }  // namespace mujoco

test/xml/xml_native_writer_test.cc

@@ -1280,8 +1280,39 @@ TEST_F(XMLWriterTest, WriteReadCompare) {
         EXPECT_EQ(d->pstack, 0) << "mjData stack memory leak detected in " <<
             p.path().string() << '\n';
 
-        // delete original structures
+        // delete data
         mj_deleteData(d);
+
+        // allocate buffer, save m into it
+        size_t sz = mj_sizeModel(m);
+        void* buffer = mju_malloc(sz);
+        mj_saveModel(m, nullptr, buffer, sz);
+
+        // make new VFS add buffer to it
+        mjVFS* vfs = (mjVFS*)mju_malloc(sizeof(mjVFS));
+        mj_defaultVFS(vfs);
+        int failed = mj_addBufferVFS(vfs, "model.mjb", buffer, sz);
+        EXPECT_EQ(failed, 0) << "Failed to add buffer to VFS";
+
+        // load model from VFS
+        mtemp = mj_loadModel("model.mjb", vfs);
+        ASSERT_THAT(mtemp, NotNull());
+
+        // compare with 0 tolerance
+        std::string field = "";
+        mjtNum result = CompareModel(m, mtemp, field);
+        EXPECT_EQ(result, 0)
+            << "Loaded and saved binary models are different!\n"
+            << "Affected file " << p.path().string() << '\n'
+            << "Different field: " << field << '\n';
+
+        // clean up
+        mj_deleteModel(mtemp);
+        mj_deleteVFS(vfs);
+        mju_free(vfs);
+        mju_free(buffer);
+
+        // delete model
         mj_deleteModel(m);
       }
     }