Add support functions for id2name and name2id.

PiperOrigin-RevId: 627525330
Change-Id: Ic8069f64c94ef4b8f96f62b88ece62a8b296d870

doc/changelog.rst

@@ -39,14 +39,15 @@ MJX
 11. Fixed a bug where capsules might be ignored in broadphase colliision checking.
 12. Added cylinder collisions using SDFs.
 13. Added support for all :ref:`condim <coContact>`: 1, 3, 4, 6.
+14. Add support functions for ``id2name`` and ``name2id``, MJX versions of :ref:`mj_id2name` and :ref:`mj_name2id`.
 
 Bug fixes
 ^^^^^^^^^
-14. Defaults of lights were not being saved, now fixed.
-15. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
-16. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
+15. Defaults of lights were not being saved, now fixed.
+16. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
+17. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
     not optional.
-16. Fixed bug that prevented memory allocations larger than 2.15 GB.
+18. Fixed bug that prevented memory allocations larger than 2.15 GB.
 
 
 Version 3.1.4 (April 10th, 2024)

mjx/mujoco/mjx/__init__.py

@@ -42,7 +42,9 @@ from mujoco.mjx._src.smooth import rne
 from mujoco.mjx._src.smooth import transmission
 from mujoco.mjx._src.solver import solve
 from mujoco.mjx._src.support import full_m
+from mujoco.mjx._src.support import id2name
 from mujoco.mjx._src.support import is_sparse
 from mujoco.mjx._src.support import mul_m
+from mujoco.mjx._src.support import name2id
 from mujoco.mjx._src.test_util import benchmark
 from mujoco.mjx._src.types import *

mjx/mujoco/mjx/_src/collision_driver.py

@@ -305,6 +305,11 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
   return groups
 
 
+def _numeric(m: Union[Model, mujoco.MjModel], name: str) -> int:
+  id_ = support.name2id(m, mujoco.mjtObj.mjOBJ_NUMERIC, name)
+  return int(m.numeric_data[id_]) if id_ >= 0 else -1
+
+
 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:
@@ -314,7 +319,8 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
 
   # max_geom_pairs limits the number of pairs we process in a collision function
   # by first running a primitive broad phase culling on the pairs
-  max_geom_pairs = support.get_custom_int(m, 'max_geom_pairs')
+  max_geom_pairs = _numeric(m, 'max_geom_pairs')
+
   if max_geom_pairs > -1:
     for k in group_counts:
       if set(k.types) & _GEOM_NO_BROADPHASE:
@@ -323,7 +329,7 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
 
   # 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')
+  max_contact_points = _numeric(m, 'max_contact_points')
 
   condim_counts = {}
   for k, v in group_counts.items():
@@ -344,8 +350,8 @@ def collision(m: Model, d: Data) -> Data:
     return d
 
   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')
+  max_geom_pairs = _numeric(m, 'max_geom_pairs')
+  max_contact_points = _numeric(m, 'max_contact_points')
 
   # run collision functions on groups
   for key, contact in groups.items():

mjx/mujoco/mjx/_src/support.py

@@ -25,6 +25,7 @@ from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import JacobianType
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
+import numpy as np
 
 
 def is_sparse(m: Union[mujoco.MjModel, Model]) -> bool:
@@ -187,16 +188,92 @@ def local_to_global(
   return pos, mat
 
 
-def get_custom_numeric(m: Union[Model, mujoco.MjModel], name: str) -> float:
-  """Returns a custom numeric given an MjModel or mjx.Model."""
-  for i in range(m.nnumeric):
-    name_ = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0]
-    if name_ == name:
-      return m.numeric_data[m.numeric_adr[i]]
-
-  return -1
+def _getnum(m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj) -> int:
+  """Gets the number of objects for the given object type."""
+  return {
+      mujoco.mjtObj.mjOBJ_BODY: m.nbody,
+      mujoco.mjtObj.mjOBJ_JOINT: m.njnt,
+      mujoco.mjtObj.mjOBJ_GEOM: m.ngeom,
+      mujoco.mjtObj.mjOBJ_SITE: m.nsite,
+      mujoco.mjtObj.mjOBJ_CAMERA: m.ncam,
+      mujoco.mjtObj.mjOBJ_MESH: m.nmesh,
+      mujoco.mjtObj.mjOBJ_PAIR: m.npair,
+      mujoco.mjtObj.mjOBJ_EQUALITY: m.neq,
+      mujoco.mjtObj.mjOBJ_ACTUATOR: m.nu,
+      mujoco.mjtObj.mjOBJ_SENSOR: m.nsensor,
+      mujoco.mjtObj.mjOBJ_NUMERIC: m.nnumeric,
+      mujoco.mjtObj.mjOBJ_TUPLE: m.ntuple,
+      mujoco.mjtObj.mjOBJ_KEY: m.nkey,
+  }.get(obj, 0)
 
 
-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))
+def _getadr(
+    m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj
+) -> np.ndarray:
+  """Gets the name addresses for the given object type."""
+  return {
+      mujoco.mjtObj.mjOBJ_BODY: m.name_bodyadr,
+      mujoco.mjtObj.mjOBJ_JOINT: m.name_jntadr,
+      mujoco.mjtObj.mjOBJ_GEOM: m.name_geomadr,
+      mujoco.mjtObj.mjOBJ_SITE: m.name_siteadr,
+      mujoco.mjtObj.mjOBJ_CAMERA: m.name_camadr,
+      mujoco.mjtObj.mjOBJ_MESH: m.name_meshadr,
+      mujoco.mjtObj.mjOBJ_PAIR: m.name_pairadr,
+      mujoco.mjtObj.mjOBJ_EQUALITY: m.name_eqadr,
+      mujoco.mjtObj.mjOBJ_ACTUATOR: m.name_actuatoradr,
+      mujoco.mjtObj.mjOBJ_SENSOR: m.name_sensoradr,
+      mujoco.mjtObj.mjOBJ_NUMERIC: m.name_numericadr,
+      mujoco.mjtObj.mjOBJ_TUPLE: m.name_tupleadr,
+      mujoco.mjtObj.mjOBJ_KEY: m.name_keyadr,
+  }[obj]
+
+
+def id2name(
+    m: Union[Model, mujoco.MjModel], typ: mujoco._enums.mjtObj, i: int
+) -> Optional[str]:
+  """Gets the name of an object with the specified mjtObj type and id.
+
+  See mujoco.id2name for more info.
+
+  Args:
+    m: mujoco.MjModel or mjx.Model
+    typ: mujoco.mjtObj type
+    i: the id
+
+  Returns:
+    the name string, or None if not found
+  """
+  num = _getnum(m, typ)
+  if i < 0 or i >= num:
+    return None
+
+  adr = _getadr(m, typ)
+  name = m.names[adr[i] :].decode('utf-8').split('\x00', 1)[0]
+  return name or None
+
+
+def name2id(
+    m: Union[Model, mujoco.MjModel], typ: mujoco._enums.mjtObj, name: str
+) -> int:
+  """Gets the id of an object with the specified mjtObj type and name.
+
+  See mujoco.mj_name2id for more info.
+
+  Args:
+    m: mujoco.MjModel or mjx.Model
+    typ: mujoco.mjtObj type
+    name: the name of the object
+
+  Returns:
+   the id, or -1 if not found
+  """
+  num = _getnum(m, typ)
+  adr = _getadr(m, typ)
+
+  # TODO: consider using MjModel.names_map instead
+  names_map = {
+      m.names[adr[i] :].decode('utf-8').split('\x00', 1)[0]: i
+      for i in range(num)
+  }
+
+  return names_map.get(name, -1)

mjx/mujoco/mjx/_src/support_test.py

@@ -116,7 +116,7 @@ class SupportTest(parameterized.TestCase):
 
     np.testing.assert_almost_equal(qfrc, qfrc_expected, 6)
 
-  def test_custom_numeric(self):
+  def test_custom(self):
     xml = """
     <mujoco model="right_shadow_hand">
         <custom>
@@ -126,9 +126,36 @@ class SupportTest(parameterized.TestCase):
     </mujoco>
     """
     m = mujoco.MjModel.from_xml_string(xml)
-    self.assertEqual(support.get_custom_numeric(m, 'something'), -1)
-    self.assertEqual(support.get_custom_numeric(m, 'max_contact_points'), 15)
-    self.assertEqual(support.get_custom_numeric(m, 'max_geom_pairs'), 42)
+
+    def _get_numeric(m, name):
+      id_ = support.name2id(m, mujoco.mjtObj.mjOBJ_NUMERIC, name)
+      return int(m.numeric_data[id_]) if id_ >= 0 else -1
+
+    self.assertEqual(_get_numeric(m, 'something'), -1)
+    self.assertEqual(_get_numeric(m, 'max_contact_points'), 15)
+    self.assertEqual(_get_numeric(m, 'max_geom_pairs'), 42)
+
+    mx = mjx.put_model(m)
+    self.assertEqual(_get_numeric(mx, 'something'), -1)
+    self.assertEqual(_get_numeric(mx, 'max_contact_points'), 15)
+    self.assertEqual(_get_numeric(mx, 'max_geom_pairs'), 42)
+
+  def test_names_and_ids(self):
+    m = test_util.load_test_file('pendula.xml')
+    mx = mjx.put_model(m)
+
+    nums = {
+        mujoco.mjtObj.mjOBJ_JOINT: m.njnt,
+        mujoco.mjtObj.mjOBJ_GEOM: m.ngeom,
+        mujoco.mjtObj.mjOBJ_BODY: m.nbody,
+    }
+
+    for obj in nums:
+      names = [mujoco.mj_id2name(m, obj.value, i) for i in range(nums[obj])]
+      for i, n in enumerate(names):
+        self.assertEqual(support.id2name(mx, obj, i), n)
+        i = i if n is not None else -1
+        self.assertEqual(support.name2id(mx, obj, n), i)
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/types.py

@@ -98,6 +98,7 @@ class GeomType(enum.IntEnum):
     CYLINDER: cylinder
     BOX: box
     MESH: mesh
+    SDF: signed distance field
   """
 
   PLANE = mujoco.mjtGeom.mjGEOM_PLANE
@@ -326,6 +327,9 @@ class Model(PyTreeNode):
     nexclude: number of excluded geom pairs
     neq: number of equality constraints
     nnumeric: number of numeric custom fields
+    ntuple: number of tuple custom fields
+    nsensor: number of sensors
+    nkey: number of keyframes
     nuserdata: size of userdata array
     nM: number of non-zeros in sparse inertia matrix
     opt: physics options
@@ -448,8 +452,24 @@ 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,)
+    tuple_adr: address of text in text_data                   (ntuple,)
+    tuple_size: number of objects in tuple                    (ntuple,)
+    tuple_objtype: array of object types in all tuples        (ntupledata,)
+    tuple_objid: array of object ids in all tuples            (ntupledata,)
+    tuple_objprm: array of object params in all tuples        (ntupledata,)
+    name_bodyadr: body name pointers                          (nbody,)
+    name_jntadr: joint name pointers                          (njnt,)
+    name_geomadr: geom name pointers                          (ngeom,)
+    name_siteadr: site name pointers                          (nsite,)
+    name_camadr: camera name pointers                         (ncam,)
     name_meshadr: mesh name pointers                          (nmesh,)
+    name_pairadr: geom pair name pointers                     (npair,)
+    name_eqadr: equality constraint name pointers             (neq,)
+    name_actuatoradr: actuator name pointers                  (nu,)
+    name_sensoradr: sensor name pointers                      (nsensor,)
     name_numericadr: numeric name pointers                    (nnumeric,)
+    name_tupleadr: tuple name pointers                        (ntuple,)
+    name_keyadr: keyframe name pointers                       (nkey,)
     names: names of all objects, 0-terminated                 (nnames,)
   """
   nq: int
@@ -470,6 +490,9 @@ class Model(PyTreeNode):
   neq: int
   nnumeric: int
   nuserdata: int
+  ntuple: int
+  nsensor: int
+  nkey: int
   nM: int  # pylint:disable=invalid-name
   opt: Option
   stat: Statistic
@@ -591,8 +614,24 @@ class Model(PyTreeNode):
   actuator_gear: jax.Array
   numeric_adr: np.ndarray
   numeric_data: np.ndarray
+  tuple_adr: np.ndarray
+  tuple_size: np.ndarray
+  tuple_objtype: np.ndarray
+  tuple_objid: np.ndarray
+  tuple_objprm: np.ndarray
+  name_bodyadr: np.ndarray
+  name_jntadr: np.ndarray
+  name_geomadr: np.ndarray
+  name_siteadr: np.ndarray
+  name_camadr: np.ndarray
   name_meshadr: np.ndarray
+  name_pairadr: np.ndarray
+  name_eqadr: np.ndarray
+  name_actuatoradr: np.ndarray
+  name_sensoradr: np.ndarray
   name_numericadr: np.ndarray
+  name_tupleadr: np.ndarray
+  name_keyadr: np.ndarray
   names: bytes