Add put_model, put_data, and get_data to MJX.

These io functions replace device_put, device_get_into. These new functions correctly copy intermediate fields such as efc_J, which the old functions could not.

PiperOrigin-RevId: 586712273
Change-Id: Ia3a4104af17db2ebd41b41a2ca501c5de9ac5d05

doc/changelog.rst

@@ -27,6 +27,9 @@ MJX
 ^^^
 
 - Added ``site_xpos`` and ``site_xmat`` to MJX.
+- Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be
+  deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as
+  ``efc_J``.
 
 Bug fixes
 ^^^^^^^^^

mjx/mujoco/mjx/__init__.py

@@ -21,7 +21,10 @@ from mujoco.mjx._src.device import device_get_into
 from mujoco.mjx._src.device import device_put
 from mujoco.mjx._src.forward import forward
 from mujoco.mjx._src.forward import step
+from mujoco.mjx._src.io import get_data
 from mujoco.mjx._src.io import make_data
+from mujoco.mjx._src.io import put_data
+from mujoco.mjx._src.io import put_model
 from mujoco.mjx._src.passive import passive
 from mujoco.mjx._src.smooth import com_pos
 from mujoco.mjx._src.smooth import com_vel

mjx/mujoco/mjx/_src/collision_driver.py

@@ -40,7 +40,6 @@ 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
@@ -81,6 +80,12 @@ def _add_candidate(
   if t1 > t2:
     t1, t2, g1, g2 = t2, t1, g2, g1
 
+  # MuJoCo does not collide planes with other planes or hfields
+  if t1 == GeomType.PLANE and t2 == GeomType.PLANE:
+    return
+  if t1 == GeomType.PLANE and t2 == GeomType.HFIELD:
+    return
+
   def mesh_key(i):
     convex_data = [[None] * m.ngeom] * 3
     if isinstance(m, Model):
@@ -284,13 +289,11 @@ def _collide_geoms(
       solimp=params.solimp,
       geom1=geom1,
       geom2=geom2,
-      dim=np.array([]),
-      efc_address=np.array([]),
   )
   return con
 
 
-def _max_contact_points(m: Model) -> int:
+def _max_contact_points(m: Union[Model, mujoco.MjModel]) -> int:
   """Returns the maximum number of contact points when set as a numeric."""
   for i in range(m.nnumeric):
     name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0]
@@ -334,7 +337,7 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
   return candidate_set
 
 
-def ncon(m: Model) -> int:
+def ncon(m: Union[Model, mujoco.MjModel]) -> int:
   """Returns the number of contacts computed in MJX given a model."""
   if m.opt.disableflags & DisableBit.CONTACT:
     return 0
@@ -354,9 +357,8 @@ def ncon(m: Model) -> int:
 
 def collision(m: Model, d: Data) -> Data:
   """Collides geometries."""
-  ncon_ = ncon(m)
+  if ncon(m) == 0:
-  if ncon_ == 0:
+    return d.replace(contact=Contact.zero())
-    return d.replace(contact=Contact.zero(), ncon=0)
 
   candidate_set = collision_candidates(m)
 
@@ -376,13 +378,4 @@ def collision(m: Model, d: Data) -> Data:
     _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points)
     contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact)
 
-  if ncon_ != contact.dist.shape[0]:
+  return d.replace(contact=contact)
-    raise RuntimeError('Number of contacts does not match ncon.')
-
-  # TODO(robotics-simulation): move this logic to device_put
-  ns = d.ne + d.nf + d.nl
-  contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4))
-  # TODO(robotics-simulation): add support for other friction dimensions
-  contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32))
-
-  return d.replace(contact=contact, ncon=ncon_)

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -253,7 +253,6 @@ class CapsuleCollisionTest(parameterized.TestCase):
     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 = c.replace(dim=c.dim[np.array([0])])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
 
@@ -281,7 +280,6 @@ class ConvexTest(absltest.TestCase):
     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 = c.replace(dim=c.dim[np.array([0, 1])])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2)
 
@@ -339,7 +337,6 @@ class ConvexTest(absltest.TestCase):
     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 = c.replace(dim=c.dim[np.array([0])])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
 
@@ -517,8 +514,8 @@ class TopKContactTest(absltest.TestCase):
     dx_all = collision_jit_fn(mx_all, dx)
     dx_top_k = collision_jit_fn(mx_top_k, dx)
 
-    self.assertEqual(dx_all.ncon, 3)
+    self.assertEqual(dx_all.contact.dist.shape, (3,))
-    self.assertEqual(dx_top_k.ncon, 2)
+    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/constraint.py

@@ -14,11 +14,12 @@
 # ==============================================================================
 """Core non-smooth constraint functions."""
 
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import jax
 from jax import numpy as jp
 import mujoco
+from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import math
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
@@ -276,7 +277,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 (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0:
+  if collision_driver.ncon(m) == 0:
     return None
 
   @jax.vmap
@@ -313,7 +314,9 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
   return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
 
 
-def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
+def count_constraints(
+    m: Union[Model, mujoco.MjModel]
+) -> Tuple[int, int, int, int]:
   """Returns equality, friction, limit, and contact constraint counts."""
   if m.opt.disableflags & DisableBit.CONSTRAINT:
     return 0, 0, 0, 0
@@ -333,10 +336,7 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
   else:
     nl = int(m.jnt_limited.sum())
 
-  if m.opt.disableflags & DisableBit.CONTACT:
+  nc = collision_driver.ncon(m) * 4
-    nc = 0
-  else:
-    nc = d.ncon * 4
 
   return ne, nf, nl, nc
 
@@ -344,10 +344,6 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
 def make_constraint(m: Model, d: Data) -> Data:
   """Creates constraint jacobians and other supporting data."""
 
-  ns = sum(count_constraints(m, d)[:-1])
-  # TODO(robotics-simulation): make device_put set nefc/efc_address instead
-  d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)})
-
   if m.opt.disableflags & DisableBit.CONSTRAINT:
     efcs = ()
   else:
@@ -364,7 +360,7 @@ def make_constraint(m: Model, d: Data) -> Data:
   if not efcs:
     z = jp.empty(0)
     d = d.replace(efc_J=jp.empty((0, m.nv)))
-    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, nefc=0)
+    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z)
     return d
 
   efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
@@ -378,6 +374,6 @@ def make_constraint(m: Model, d: Data) -> Data:
 
   aref, r = fn(efc)
   d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
-  d = d.replace(efc_frictionloss=efc.frictionloss, nefc=r.shape[0])
+  d = d.replace(efc_frictionloss=efc.frictionloss)
 
   return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -136,13 +136,8 @@ class ConstraintTest(parameterized.TestCase):
     m = test_util.load_test_file('ant.xml')
     d = mujoco.MjData(m)
 
-    m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
-    dx = constraint.make_constraint(mx, dx)
-    self.assertGreater(dx.efc_J.shape[0], 1)
-
     m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT
-    mx = mjx.device_put(m)
+    mx, dx = mjx.device_put(m), mjx.device_put(d)
     dx = constraint.make_constraint(mx, dx)
     self.assertEqual(dx.efc_J.shape[0], 0)
 

mjx/mujoco/mjx/_src/dataclasses.py

@@ -18,7 +18,7 @@ import copy
 import dataclasses
 
 import typing
-from typing import Dict, Optional, Sequence, TypeVar
+from typing import Any, Dict, Optional, Sequence, TypeVar
 import jax
 import numpy as np
 
@@ -113,6 +113,10 @@ class PyTreeNode:
     # stub for pytype
     raise NotImplementedError
 
+  @classmethod
+  def fields(cls) -> tuple[dataclasses.Field[Any], ...]:
+    return dataclasses.fields(cls)
+
   def tree_replace(
       self, params: Dict[str, Optional[jax.typing.ArrayLike]]
   ) -> 'PyTreeNode':

mjx/mujoco/mjx/_src/device.py

@@ -25,6 +25,7 @@ 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,),
@@ -68,7 +69,6 @@ _TRANSFORMS = {
     (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.Model, 'actuator_trnid'): lambda x: x[:, 0],
     (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))
@@ -265,9 +265,16 @@ def device_get_into(result, value):
 
   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=value.ncon, nefc=value.nefc
+          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:

mjx/mujoco/mjx/_src/io.py

@@ -14,94 +14,337 @@
 # ==============================================================================
 """Functions to initialize, load, or save data."""
 
+import copy
+from typing import List, Union
+
+import jax
 from jax import numpy as jp
+import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
-# pylint: disable=g-importing-member
+from mujoco.mjx._src import mesh
-from mujoco.mjx._src.types import Contact
+from mujoco.mjx._src import types
-from mujoco.mjx._src.types import Data
-from mujoco.mjx._src.types import Model
-# pylint: enable=g-importing-member
 import numpy as np
 
 
-def make_data(m: Model) -> Data:
+def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
+  """Puts mujoco.MjOption onto a device, resulting in mjx.Option."""
+  if o.integrator not in set(types.IntegratorType):
+    raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
+
+  if o.cone not in set(types.ConeType):
+    raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}')
+
+  if o.solver not in set(types.SolverType):
+    raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
+
+  for i in range(mujoco.mjtEnableBit.mjNENABLE):
+    if o.enableflags & 2**i:
+      raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
+
+  static_fields = {
+      f.name: copy.copy(getattr(o, f.name))
+      for f in types.Option.fields()
+      if f.type in (int, bytes, np.ndarray)
+  }
+  static_fields['integrator'] = types.IntegratorType(o.integrator)
+  static_fields['cone'] = types.ConeType(o.cone)
+  static_fields['solver'] = types.SolverType(o.solver)
+  static_fields['disableflags'] = types.DisableBit(o.disableflags)
+
+  device_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:
+  """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic."""
+  return types.Statistic(
+      meaninertia=jax.device_put(s.meaninertia, device=device)
+  )
+
+
+def put_model(m: mujoco.MjModel, device=None) -> types.Model:
+  """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
+  if m.ntendon:
+    raise NotImplementedError('tendons are not supported')
+
+  if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
+    raise NotImplementedError('only condim=3 is supported')
+
+  # check collision geom types
+  for g1, g2, *_ in collision_driver.collision_candidates(m):
+    if collision_driver.get_collision_fn((g1, g2)) is None:
+      g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
+      raise NotImplementedError(f'({g1}, {g2}) has no collision function')
+
+  for enum_field, enum_type, mj_type in (
+      (m.actuator_biastype, types.BiasType, mujoco.mjtBias),
+      (m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
+      (m.actuator_gaintype, types.GainType, mujoco.mjtGain),
+      (m.actuator_trntype, types.TrnType, mujoco.mjtTrn),
+      (m.eq_type, types.EqType, mujoco.mjtEq),
+  ):
+    missing = set(enum_field) - set(enum_type)
+    if missing:
+      raise NotImplementedError(
+          f'{[mj_type(m) for m in missing]} not supported'
+      )
+
+  opt = _put_option(m.opt, device=device)
+  stat = _put_statistic(m.stat, device=device)
+
+  static_fields = {
+      f.name: getattr(m, f.name)
+      for f in types.Model.fields()
+      if f.type in (int, bytes, np.ndarray)
+  }
+
+  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.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."""
 
+  ncon = collision_driver.ncon(m)
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  nefc = ne + nf + nl + nc
+
+  zero_nv = jp.zeros(m.nv, dtype=jp.float32)
+  zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32)
+  zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32)
+  zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32)
+  zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32)
+  zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32)
+  zero_nefc = jp.zeros(nefc, dtype=jp.float32)
+  zero_na = jp.zeros(m.na, dtype=jp.float32)
+  zero_nu = jp.zeros(m.nu, dtype=jp.float32)
+  zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32)
+  zero_nm = jp.zeros(m.nM, dtype=jp.float32)
+
   # create first d to get num contacts and nc
-  d = Data(
+  d = types.Data(
       solver_niter=jp.array(0, dtype=jp.int32),
-      ne=0,
+      time=jp.array(0.0),
-      nf=0,
-      nl=0,
-      nefc=0,
-      ncon=0,
-      time=jp.zeros((), dtype=jp.float32),
       qpos=m.qpos0,
-      qvel=jp.zeros(m.nv, dtype=jp.float32),
+      qvel=zero_nv,
-      act=jp.zeros(m.na, dtype=jp.float32),
+      act=zero_na,
-      qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32),
+      qacc_warmstart=zero_nv,
-      ctrl=jp.zeros(m.nu, dtype=jp.float32),
+      ctrl=zero_nu,
-      qfrc_applied=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_applied=zero_nv,
-      xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32),
+      xfrc_applied=zero_nbody_6,
       eq_active=jp.zeros(m.neq, dtype=jp.int32),
-      qacc=jp.zeros(m.nv, dtype=jp.float32),
+      qacc=zero_nv,
-      act_dot=jp.zeros(m.na, dtype=jp.float32),
+      act_dot=zero_na,
-      xpos=jp.zeros((m.nbody, 3), dtype=jp.float32),
+      xpos=zero_nbody_3,
       xquat=jp.zeros((m.nbody, 4), dtype=jp.float32),
-      xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32),
+      xmat=zero_nbody_3_3,
-      xipos=jp.zeros((m.nbody, 3), dtype=jp.float32),
+      xipos=zero_nbody_3,
-      ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32),
+      ximat=zero_nbody_3_3,
-      xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32),
+      xanchor=zero_njnt_3,
-      xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32),
+      xaxis=zero_njnt_3,
       geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
       geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
       site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
       site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
-      subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32),
+      subtree_com=zero_nbody_3,
-      cdof=jp.zeros((m.nv, 6), dtype=jp.float32),
+      cdof=zero_nv_6,
-      cinert=jp.zeros((m.nbody, 10), dtype=jp.float32),
+      cinert=zero_nbody_10,
-      actuator_length=jp.zeros(m.nu, dtype=jp.float32),
+      actuator_length=zero_nu,
       actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
-      crb=jp.zeros((m.nbody, 10), dtype=jp.float32),
+      crb=zero_nbody_10,
-      qM=jp.zeros(m.nM, dtype=jp.float32),
+      qM=zero_nm,
-      qLD=jp.zeros(m.nM, dtype=jp.float32),
+      qLD=zero_nm,
-      qLDiagInv=jp.zeros(m.nv, dtype=jp.float32),
+      qLDiagInv=zero_nv,
-      qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32),
+      qLDiagSqrtInv=zero_nv,
-      contact=Contact.zero(),
+      contact=types.Contact.zero(ncon),
-      efc_J=jp.zeros((), dtype=jp.float32),
+      efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32),
-      efc_frictionloss=jp.zeros((), dtype=jp.float32),
+      efc_frictionloss=zero_nefc,
-      efc_D=jp.zeros((), dtype=jp.float32),
+      efc_D=zero_nefc,
-      actuator_velocity=jp.zeros(m.nu, dtype=jp.float32),
+      actuator_velocity=zero_nu,
-      cvel=jp.zeros((m.nbody, 6), dtype=jp.float32),
+      cvel=zero_nbody_6,
-      cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32),
+      cdof_dot=zero_nv_6,
-      qfrc_bias=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_bias=zero_nv,
-      qfrc_passive=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_passive=zero_nv,
-      efc_aref=jp.zeros((), dtype=jp.float32),
+      efc_aref=zero_nefc,
-      actuator_force=jp.zeros(m.nu, dtype=jp.float32),
+      actuator_force=zero_nu,
-      qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_actuator=zero_nv,
-      qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_smooth=zero_nv,
-      qacc_smooth=jp.zeros(m.nv, dtype=jp.float32),
+      qacc_smooth=zero_nv,
-      qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_constraint=zero_nv,
-      qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32),
+      qfrc_inverse=zero_nv,
-      efc_force=jp.zeros((), dtype=jp.float32),
+      efc_force=zero_nefc,
-  )
-
-  # get contact data with correct shapes
-  ncon = collision_driver.ncon(m)
-  d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon)
-  d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)})
-
-  ne, nf, nl, nc = constraint.count_constraints(m, d)
-  d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc)
-  ns = ne + nf + nl
-  d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)})
-  d = d.replace(
-      efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32),
-      efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32),
-      efc_D=jp.zeros(d.nefc, dtype=jp.float32),
-      efc_aref=jp.zeros(d.nefc, dtype=jp.float32),
-      efc_force=jp.zeros(d.nefc, dtype=jp.float32),
   )
 
   return d
+
+
+def _get_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():
+    value = getattr(cx, field.name)[con_id]
+    if field.name == 'frame':
+      value = value.reshape((-1, 9))
+    getattr(c, field.name)[:] = value
+
+  ncon = con_id.shape[0]
+  c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
+
+
+def get_data(
+    m: mujoco.MjModel, d: types.Data
+) -> Union[mujoco.MjData, List[mujoco.MjData]]:
+  """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData]."""
+  dx = jax.device_get(d)
+  batched = len(d.qpos.shape) > 1
+  batch_size = d.qpos.shape[0] if batched else 1
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  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])
+
+  ds = []
+  for i in range(batch_size):
+    dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d
+    ncon = (dx_i.contact.dist <= 0).sum()
+    efc_active = (dx_i.efc_J != 0).any(axis=1)
+    efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
+    nefc, nc = efc_active.sum(), (efc_active & efc_con).sum()
+    d_i = mujoco.MjData(m)
+    d_i.nnzJ = nefc * m.nv
+    mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc)  # pylint: disable=protected-access
+    d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
+    d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
+    d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
+
+    for field in types.Data.fields():
+      if field.name == 'contact':
+        _get_contact(d_i.contact, dx_i.contact, nefc - nc)
+        continue
+
+      value = getattr(dx_i, field.name)
+
+      if field.name in ('xmat', 'ximat', 'geom_xmat'):
+        value = value.reshape((-1, 9))
+
+      if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+        value = value[efc_active]
+
+      if field.name == 'efc_J':
+        value = value[efc_active].reshape(-1)
+
+      if value.shape:
+        getattr(d_i, field.name)[:] = value
+      else:
+        setattr(d_i, field.name, value)
+
+    d_i.efc_type[:] = efc_type[efc_active]
+    ds.append(d_i)
+
+  return ds if batched else ds[0]
+
+
+def _put_contact(
+    c: mujoco._structs._MjContactList, ncon: int, device=None
+) -> types.Contact:
+  """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.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_fn = lambda x: np.concatenate(
+      (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
+  )
+  fields = jax.tree_map(pad_fn, fields)
+  fields['dist'][-pad_size:] = np.inf
+  fields = jax.device_put(fields, device=device)
+
+  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)
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  nefc = ne + nf + nl + nc
+
+  for d_val, val, name in (
+      (d.ncon, ncon, 'ncon'),
+      (d.ne, ne, 'ne'),
+      (d.nf, nf, 'nf'),
+      (d.nl, nl, 'nl'),
+      (d.nefc, nefc, 'nefc'),
+  ):
+    if d_val > val:
+      raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
+
+  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
+  }
+
+  for fname in ('xmat', 'ximat', 'geom_xmat'):
+    fields[fname] = fields[fname].reshape((-1, 3, 3))
+
+  # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
+  # efc_J is also optionally column-sparse (typically for large nv).  MJX is
+  # neither: it contains zeros for inactive constraints, and efc_J is always
+  # (nefc, nv).  this may change in the future.
+  if mujoco.mj_isSparse(m):
+    nr = d.efc_J_rownnz.shape[0]
+    efc_j = np.zeros((nr, m.nv))
+    for i in range(nr):
+      rowadr = d.efc_J_rowadr[i]
+      for j in range(d.efc_J_rownnz[i]):
+        efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
+    fields['efc_J'] = efc_j
+  else:
+    fields['efc_J'] = fields['efc_J'].reshape((-1, m.nv))
+
+  for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+    value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc)
+    for i in range(4):
+      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]
+    fields[fname] = value
+
+  fields = jax.device_put(fields, device=device)
+  fields['contact'] = _put_contact(d.contact, ncon, device=device)
+
+  return types.Data(**fields)

mjx/mujoco/mjx/_src/io_test.py

@@ -17,26 +17,390 @@
 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 test_util
+import numpy as np
+
+
+_MULTIPLE_CONVEX_OBJECTS = """
+  <mujoco>
+    <option timestep="0.001"/>
+    <default>
+      <geom solref=".006 1"/>
+    </default>
+    <asset>
+      <mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 .1"/>
+      <mesh name="boxoid" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 .5 -1 -1 2" scale=".3 .2 .1"/>
+      <mesh name="pentaprism" vertex="1 0 0 0.309 0.951 0 -0.809 0.588 0 -0.809 -0.588 0 0.309 -0.951 0 1 0 1 0.309 0.951 1 -0.809 0.588 1 -0.809 -0.588 1 0.309 -0.951 1" scale=".2 .2 .1"/>
+    </asset>
+    <worldbody>
+      <geom type="plane" pos="0 0 -.5" size="3 3 .01"/>
+      <geom type="mesh" mesh="box" pos="0 0 -.15" euler="3 7 30"/>
+      <body pos="-.3 -.3 .3">
+        <freejoint/>
+        <geom type="mesh" mesh="boxoid" rgba=".8 0 0 1" euler="3 5 -130"/>
+      </body>
+      <body pos=".3 .3 0.3">
+        <freejoint/>
+        <geom type="box" euler="3 5 -80" size=".3 .2 .1" rgba="0 .8 0 1"/>
+      </body>
+      <body pos=".3 .3 .6">
+        <freejoint/>
+        <geom type="mesh" mesh="pentaprism" rgba="0 0 .8 1"/>
+      </body>
+      <body pos=".6 -.3 .3">
+        <joint name="joint" axis="1 0 0" type="hinge" range="-45 45"/>
+        <geom type="capsule" size=".2 .05" rgba=".6 0 .6 1"/>
+      </body>
+    </worldbody>
+    <actuator>
+      <motor joint="joint"/>
+    </actuator>
+  </mujoco>
+"""
+
+_MULTIPLE_CONSTRAINTS = """
+  <mujoco>
+    <worldbody>
+      <geom type="plane" size="3 3 .01"/>
+      <body name="cap1" pos="-.3 -.3 .2">
+        <freejoint/>
+        <geom type="capsule" size=".2 .05"/>
+        <body name="cap2" pos=".6 -.3 .3">
+          <joint axis="0 1 0" type="hinge" range="-45 45"/>
+          <joint axis="1 0 0" type="hinge" range="-0.001 0.001"/>
+          <geom type="capsule" size=".2 .05"/>
+        </body>
+      </body>
+    </worldbody>
+    <equality>
+      <connect body1="cap2" anchor="0 0 1"/>
+    </equality>
+  </mujoco>
+"""
 
 
 class IoTest(parameterized.TestCase):
 
-  @parameterized.parameters(test_util.TEST_FILES)
+  def test_put_model(self):
-  def test_make_data(self, fname):
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
-    """Test that data created by make_data matches data returned by step."""
+    mx = mjx.put_model(m)
+    self.assertEqual(mx.nq, m.nq)
+    self.assertEqual(mx.nv, m.nv)
+    self.assertEqual(mx.nu, m.nu)
+    self.assertEqual(mx.na, m.na)
+    self.assertEqual(mx.nbody, m.nbody)
+    self.assertEqual(mx.njnt, m.njnt)
+    self.assertEqual(mx.ngeom, m.ngeom)
+    self.assertEqual(mx.nmesh, m.nmesh)
+    self.assertEqual(mx.npair, m.npair)
+    self.assertEqual(mx.nexclude, m.nexclude)
+    self.assertEqual(mx.neq, m.neq)
+    self.assertEqual(mx.nnumeric, m.nnumeric)
+    self.assertEqual(mx.nM, m.nM)
+    self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
 
-    m = test_util.load_test_file(fname)
+    np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
-    mx = mjx.device_put(m)
+    np.testing.assert_allclose(mx.geom_type, m.geom_type)
-    dx = mjx.make_data(mx)
+    np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)
-    dx_step = mjx.step(mx, dx)
+    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, 6)
+    self.assertLen(mx.geom_convex_facenormal, 6)
 
-    _, dx_treedef = jax.tree_util.tree_flatten(dx)
+    np.testing.assert_allclose(mx.jnt_type, m.jnt_type)
-    _, dx_step_treedef = jax.tree_util.tree_flatten(dx_step)
+    np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr)
+    np.testing.assert_allclose(mx.jnt_bodyid, m.jnt_bodyid)
+    np.testing.assert_allclose(mx.jnt_limited, m.jnt_limited)
+    np.testing.assert_almost_equal(mx.jnt_axis, m.jnt_axis)
 
-    self.assertEqual(dx_treedef, dx_step_treedef)
+    np.testing.assert_allclose(mx.actuator_trntype, m.actuator_trntype)
+    np.testing.assert_allclose(mx.actuator_dyntype, m.actuator_dyntype)
+    np.testing.assert_allclose(mx.actuator_gaintype, m.actuator_gaintype)
+    np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype)
+    np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid)
 
+  def test_fluid_params(self):
+    """Test that has_fluid_params is set when fluid params are present."""
+    m = mjx.put_model(
+        mujoco.MjModel.from_xml_string(
+            '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
+        )
+    )
+    self.assertTrue(m.opt.has_fluid_params)
+
+  def test_put_model_implicit_not_implemented(self):
+    """Test that MJX guards against models with unimplemented features."""
+
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_put_model_cone_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_put_model_pgs_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option solver="PGS"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_put_model_site_actuator_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <site name="ref"/>
+            <body>
+              <site name="end"/>
+            </body>
+          </worldbody>
+          <actuator>
+            <general site="end" refsite="ref"/>
+          </actuator>
+        </mujoco>"""))
+
+  def test_put_model_tendon_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint name="left_hip" type="hinge"/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+          <tendon>
+            <fixed>
+              <joint coef="1" joint="left_hip"/>
+            </fixed>
+          </tendon>
+        </mujoco>"""))
+
+  def test_put_model_condim_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom size="0.05" condim="1"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05" condim="1"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
+  def test_put_model_cylinder_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom type="cylinder" size="0.05 0.05"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
+  def test_make_data(self):
+    """Test that make_data returns the correct shapes."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
+    d = mjx.make_data(m)
+
+    nq = 22
+    nbody = 5
+    ncon = 46
+    nv = 19
+    nefc = 185
+    nm = 64
+
+    self.assertEqual(d.qpos.shape, (nq,))
+    self.assertEqual(d.qvel.shape, (nv,))
+    self.assertEqual(d.act.shape, (0,))
+    self.assertEqual(d.qacc_warmstart.shape, (nv,))
+    self.assertEqual(d.ctrl.shape, (1,))
+    self.assertEqual(d.qfrc_applied.shape, (nv,))
+    self.assertEqual(d.xfrc_applied.shape, (nbody, 6))
+    self.assertEqual(d.eq_active.shape, (0,))
+    self.assertEqual(d.qacc.shape, (nv,))
+    self.assertEqual(d.act_dot.shape, (0,))
+    self.assertEqual(d.xpos.shape, (nbody, 3))
+    self.assertEqual(d.xquat.shape, (nbody, 4))
+    self.assertEqual(d.xmat.shape, (nbody, 3, 3))
+    self.assertEqual(d.xipos.shape, (nbody, 3))
+    self.assertEqual(d.ximat.shape, (nbody, 3, 3))
+    self.assertEqual(d.xanchor.shape, (4, 3))
+    self.assertEqual(d.xaxis.shape, (4, 3))
+    self.assertEqual(d.geom_xpos.shape, (6, 3))
+    self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
+    self.assertEqual(d.subtree_com.shape, (nbody, 3))
+    self.assertEqual(d.cdof.shape, (nv, 6))
+    self.assertEqual(d.cinert.shape, (nbody, 10))
+    self.assertEqual(d.crb.shape, (nbody, 10))
+    self.assertEqual(d.actuator_length.shape, (1,))
+    self.assertEqual(d.actuator_moment.shape, (1, nv))
+    self.assertEqual(d.qM.shape, (nm,))
+    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d.qLDiagInv.shape, (nv,))
+    self.assertEqual(d.qLDiagSqrtInv.shape, (nv,))
+    self.assertEqual(d.contact.dist.shape, (ncon,))
+    self.assertEqual(d.contact.pos.shape, (ncon, 3))
+    self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
+    self.assertEqual(d.contact.solref.shape, (ncon, 2))
+    self.assertEqual(d.contact.solimp.shape, (ncon, 5))
+    self.assertEqual(d.contact.geom1.shape, (ncon,))
+    self.assertEqual(d.contact.geom2.shape, (ncon,))
+    self.assertEqual(d.efc_J.shape, (nefc, nv))
+    self.assertEqual(d.efc_frictionloss.shape, (nefc,))
+    self.assertEqual(d.efc_D.shape, (nefc,))
+    self.assertEqual(d.actuator_velocity.shape, (1,))
+    self.assertEqual(d.cvel.shape, (nbody, 6))
+    self.assertEqual(d.cdof_dot.shape, (nv, 6))
+    self.assertEqual(d.qfrc_bias.shape, (nv,))
+    self.assertEqual(d.qfrc_passive.shape, (nv,))
+    self.assertEqual(d.efc_aref.shape, (nefc,))
+    self.assertEqual(d.actuator_force.shape, (1,))
+    self.assertEqual(d.qfrc_actuator.shape, (nv,))
+    self.assertEqual(d.qfrc_smooth.shape, (nv,))
+    self.assertEqual(d.qacc_smooth.shape, (nv,))
+    self.assertEqual(d.qfrc_constraint.shape, (nv,))
+    self.assertEqual(d.qfrc_inverse.shape, (nv,))
+    self.assertEqual(d.efc_force.shape, (nefc,))
+
+  def test_put_data(self):
+    """Test that put_data puts the correct data for dense and sparse."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+
+    # check a few fields
+    np.testing.assert_allclose(dx.qpos, d.qpos)
+    np.testing.assert_allclose(dx.xpos, d.xpos)
+    np.testing.assert_allclose(dx.cvel, d.cvel)
+    np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(dx.qM, d.qM)
+
+    # 4 contacts, 2 for each capsule against the plane
+    self.assertEqual(dx.contact.dist.shape, (4,))
+    self.assertEqual(d.ncon, 1)  # however only 1 contact in this step
+    np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
+    self.assertTrue(np.isinf(dx.contact.dist[1:]).all())
+    self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
+    np.testing.assert_allclose(
+        dx.contact.frame[0].reshape(9), d.contact.frame[0]
+    )
+    np.testing.assert_allclose(dx.contact.frame[1:], 0)
+
+    # xmat, ximat, geom_xmat are all shape transformed
+    self.assertEqual(dx.xmat.shape, (3, 3, 3))
+    self.assertEqual(dx.ximat.shape, (3, 3, 3))
+    self.assertEqual(dx.geom_xmat.shape, (3, 3, 3))
+    np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat)
+    np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat)
+    np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat)
+
+    # efc_ are also shape transformed and padded
+    self.assertEqual(dx.efc_J.shape, (21, 8))  # nefc, nv
+    d_efc_j = d.efc_J.reshape((-1, 8))
+    np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3])  # connect eq
+    np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3])  # one active limit
+    np.testing.assert_allclose(dx.efc_J[4], 0)  # one inactive limit
+    np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8])  # contact
+    np.testing.assert_allclose(dx.efc_J[9:], 0)  # no contact
+
+    # check another efc_ too
+    self.assertEqual(dx.efc_aref.shape, (21,))  # nefc
+    np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
+    np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
+    np.testing.assert_allclose(dx.efc_aref[4], 0)
+    np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8])
+    np.testing.assert_allclose(dx.efc_aref[9:], 0)
+
+    # check sparse transform is correct
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx_from_sparse = mjx.put_data(m, d)
+    np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
+
+  def test_get_data(self):
+    """Test that get_data makes correct MjData."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    d_2: mujoco.MjData = mjx.get_data(m, dx)
+
+    # check a few fields
+    np.testing.assert_allclose(d_2.qpos, d.qpos)
+    np.testing.assert_allclose(d_2.xpos, d.xpos)
+    np.testing.assert_allclose(d_2.cvel, d.cvel)
+    np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(d_2.qM, d.qM)
+
+    # only 1 contact active
+    self.assertEqual(d_2.contact.dist.shape, (1,))
+    self.assertEqual(d_2.ncon, 1)
+    np.testing.assert_allclose(d_2.contact.dist, d.contact.dist)
+    self.assertEqual(d_2.contact.frame.shape, (1, 9))
+    np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
+
+    # xmat, ximat, geom_xmat are all shape transformed
+    self.assertEqual(d_2.xmat.shape, (3, 9))
+    self.assertEqual(d_2.ximat.shape, (3, 9))
+    self.assertEqual(d_2.geom_xmat.shape, (3, 9))
+    np.testing.assert_allclose(d_2.xmat, d.xmat)
+    np.testing.assert_allclose(d_2.ximat, d.ximat)
+    np.testing.assert_allclose(d_2.geom_xmat, d.geom_xmat)
+
+    # efc_* are also shape transformed and filtered
+    self.assertEqual(d_2.efc_J.shape, (64,))  # nefc * nv
+    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):
+    """Test that get_data makes correct List[MjData] for batched Data."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    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)
+    ds = mjx.get_data(m, dx_b)
+    self.assertLen(ds, 2)
+    np.testing.assert_allclose(ds[0].qpos, d.qpos)
+    np.testing.assert_allclose(ds[1].qpos, d.qpos + 0.05, atol=1e-8)
+    self.assertEqual(ds[0].ncon, 1)
+    self.assertEqual(ds[1].ncon, 0)
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/scan.py

@@ -221,7 +221,7 @@ def flat(
         'u': i,
         'a': m.actuator_actadr[i],
         'j': (
-            m.actuator_trnid[i]
+            m.actuator_trnid[i, 0]
             if m.actuator_trntype[i] == TrnType.JOINT
             else np.array(-1)
         ),

mjx/mujoco/mjx/_src/scan_test.py

@@ -210,15 +210,16 @@ class ScanTest(absltest.TestCase):
         m, fn, 'ujqva', 'ujqva', *args, group_by='u'
     )
 
+    actuator_trnid = m.actuator_trnid[:, 0]
     np.testing.assert_array_equal(gear, m.actuator_gear)
-    np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid])
+    np.testing.assert_array_equal(jnt_typ, m.jnt_type[actuator_trnid])
     np.testing.assert_array_equal(act, jp.array([1.4, 1.1]))
     expected_vadr = np.concatenate(
-        [np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid]
+        [np.nonzero(m.dof_jntid == trnid)[0] for trnid in actuator_trnid]
     )
     np.testing.assert_array_equal(vadr, expected_vadr)
     expected_qadr = np.concatenate(
-        [np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid]
+        [np.nonzero(scan._q_jointid(m) == i)[0] for i in actuator_trnid]
     )
     np.testing.assert_array_equal(qadr, expected_qadr)
 

mjx/mujoco/mjx/_src/solver.py

@@ -19,6 +19,7 @@ from typing import Optional
 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
 # pylint: disable=g-importing-member
@@ -69,12 +70,12 @@ class _Context(PyTreeNode):
     # TODO(robotics-team): determine nv at which sparse mul is faster
     M = smooth.dense_m(m, d) if m.nv < 100 else None  # pylint: disable=invalid-name
     ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
-    nv_0 = jp.zeros((m.nv,))
+    nv_0 = jp.zeros(m.nv)
     ctx = _Context(
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
-        efc_force=jp.zeros(d.nefc),
+        efc_force=d.efc_force,
         M=M,
         Ma=ma,
         grad=nv_0,
@@ -111,7 +112,7 @@ class _LSPoint(PyTreeNode):
   @classmethod
   def create(
       cls,
-      d: Data,
+      m: Model,
       ctx: _Context,
       alpha: jax.Array,
       jv: jax.Array,
@@ -122,7 +123,8 @@ class _LSPoint(PyTreeNode):
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
 
     # TODO(robotics-team): change this to support friction constraints
-    active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
+    ne, nf, *_ = constraint.count_constraints(m)
+    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)
 
@@ -177,12 +179,11 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   Returns:
     context with new constraint force and costs
   """
-  del m
-
   # TODO(robotics-team): add friction constraints
 
   # only count active constraints
-  active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
+  ne, nf, *_ = constraint.count_constraints(m)
+  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
@@ -221,7 +222,8 @@ 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:
-    active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
+    ne, nf, *_ = constraint.count_constraints(m)
+    active = (ctx.Jaref < 0).at[:ne + nf].set(True)
     h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
     h = smooth.dense_m(m, d) + h
     h_ = jax.scipy.linalg.cho_factor(h)
@@ -265,7 +267,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 alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss)
+  point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss)
 
   def cond(ctx: _LSContext) -> jax.Array:
     done = ctx.ls_iter >= m.opt.ls_iterations

mjx/mujoco/mjx/_src/types.py

@@ -21,6 +21,7 @@ import jax
 import jax.numpy as jp
 import mujoco
 # pylint: disable=g-importing-member
+from mujoco.mjx._src import dataclasses
 from mujoco.mjx._src.dataclasses import PyTreeNode
 # pylint: enable=g-importing-member
 import numpy as np
@@ -213,7 +214,6 @@ class Option(PyTreeNode):
     integrator:       integration mode
     cone:             type of friction cone
     solver:           solver algorithm
-    integrator:       integration mode
     iterations:       number of main solver iterations
     ls_iterations:    maximum number of CG/Newton linesearch iterations
     disableflags:     bit flags for disabling standard features
@@ -329,8 +329,8 @@ class Model(PyTreeNode):
     geom_margin: include in solver if dist<margin-gap         (ngeom,)
     geom_gap: include in solver if dist<margin-gap            (ngeom,)
     site_bodyid: id of site's body                            (nsite,)
-    site_pos: local position offset rel. to body              (nsite,3)
+    site_pos: local position offset rel. to body              (nsite, 3)
-    site_quat: local orientation offset rel. to body          (nsite,4)
+    site_quat: local orientation offset rel. to body          (nsite, 4)
     geom_convex_face: vertex face data, MJX only              (ngeom,)
     geom_convex_vert: vertex data, MJX only                   (ngeom,)
     geom_convex_edge: unique edge data, MJX only              (ngeom,)
@@ -508,10 +508,8 @@ class Contact(PyTreeNode):
     solref: constraint solver reference, normal direction             (mjNREF,)
     solreffriction: constraint solver reference, friction directions  (mjNREF,)
     solimp: constraint solver impedance                               (mjNIMP,)
-    dim: contact space dimensionality: 1, 3, 4 or 6
     geom1: id of geom 1
     geom2: id of geom 2
-    efc_address: address in efc; -1: not included
   """
   dist: jax.Array
   pos: jax.Array
@@ -521,29 +519,25 @@ class Contact(PyTreeNode):
   solref: jax.Array
   solreffriction: jax.Array
   solimp: jax.Array
-  # unsupported: mu, H
+  # unsupported: mu, H, dim
-  dim: np.ndarray
   geom1: jax.Array
   geom2: jax.Array
-  efc_address: np.ndarray
+  # unsupported: efc_address, exclude
-  # unsupported: exclude
 
   @classmethod
-  def zero(cls, shape=(0,)) -> 'Contact':
+  def zero(cls, ncon: int = 0) -> 'Contact':
     """Returns a contact filled with zeros."""
     return Contact(
-        dist=jp.zeros(shape),
+        dist=jp.zeros(ncon),
-        pos=jp.zeros(shape + (3,)),
+        pos=jp.zeros((ncon, 3,)),
-        frame=jp.zeros(shape + (3, 3)),
+        frame=jp.zeros((ncon, 3, 3)),
-        includemargin=jp.zeros(shape),
+        includemargin=jp.zeros(ncon),
-        friction=jp.zeros(shape + (5,)),
+        friction=jp.zeros((ncon, 5)),
-        solref=jp.zeros(shape + (mujoco.mjNREF,)),
+        solref=jp.zeros((ncon, mujoco.mjNREF)),
-        solreffriction=jp.zeros(shape + (mujoco.mjNREF,)),
+        solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
-        solimp=jp.zeros(shape + (mujoco.mjNIMP,)),
+        solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
-        dim=np.zeros(shape, dtype=np.int32),
+        geom1=jp.zeros(ncon, dtype=jp.int32),
-        geom1=jp.zeros(shape, dtype=jp.int32),
+        geom2=jp.zeros(ncon, dtype=jp.int32),
-        geom2=jp.zeros(shape, dtype=jp.int32),
-        efc_address=np.zeros(shape, dtype=np.int32),
     )
 
 
@@ -552,11 +546,6 @@ class Data(PyTreeNode):
 
   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: nubmer of contacts
     time: simulation time
     qpos: position                                                (nq,)
     qvel: velocity                                                (nv,)
@@ -610,12 +599,6 @@ class Data(PyTreeNode):
   """
   # solver statistics:
   solver_niter: jax.Array
-  # sizes (variable in MJ, constant in MJX)
-  ne: int
-  nf: int
-  nl: int
-  nefc: int
-  ncon: int
   # global properties:
   time: jax.Array
   # state:

mjx/mujoco/mjx/integration_test/collision_driver_test.py

@@ -82,7 +82,6 @@ class CollisionDriverIntegrationTest(parameterized.TestCase):
     mjx_contact = jax.tree_map(
         lambda x: x.take(np.array(idx), axis=0), dx.contact
     )
-    mjx_contact = mjx_contact.replace(dim=mjx_contact.dim[idx])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)