chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support functions for id2name and name2id.

Browse Source 
PiperOrigin-RevId: 627525330
Change-Id: Ic8069f64c94ef4b8f96f62b88ece62a8b296d870
This commit is contained in:
Baruch Tabanpour
2024-04-23 15:34:46 -07:00
committed by Copybara-Service
parent cf433a4ad9
commit d15db54537
6 changed files with 175 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_driver.py
+10 -4
View File
@@ -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
+88 -11
View File
@@ -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
+31 -4
View File
@@ -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
+39
View File
@@ -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