From d15db54537189ee9ec382d51b9719184ceeec7fb Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 23 Apr 2024 15:34:46 -0700 Subject: [PATCH] Add support functions for id2name and name2id. PiperOrigin-RevId: 627525330 Change-Id: Ic8069f64c94ef4b8f96f62b88ece62a8b296d870 --- doc/changelog.rst | 9 ++- mjx/mujoco/mjx/__init__.py | 2 + mjx/mujoco/mjx/_src/collision_driver.py | 14 +++- mjx/mujoco/mjx/_src/support.py | 99 ++++++++++++++++++++++--- mjx/mujoco/mjx/_src/support_test.py | 35 ++++++++- mjx/mujoco/mjx/_src/types.py | 39 ++++++++++ 6 files changed, 175 insertions(+), 23 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index f931baed..d705b30d 100644 --- a/doc/changelog.rst +++ b/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 `: 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) diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 120db360..0d59c95d 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/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 * diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 7dca95e2..b67168d3 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/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(): diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index d3ada834..7b17fcfa 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 76606ca7..4aaebad4 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/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 = """ @@ -126,9 +126,36 @@ class SupportTest(parameterized.TestCase): """ 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__': diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 3986d064..65f8e356 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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