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

Cast mjx dict/set lookup enums to int for pybind11 v1.18 compatibility.

Browse Source 
PiperOrigin-RevId: 966094569
Change-Id: I014226016df9d519f250d0fc25045064cd94cb1e
This commit is contained in:
Michael Moss
2026-08-17 11:47:15 -07:00
committed by Copybara-Service
parent 8847823721
commit 06aa6c075f
4 changed files with 39 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_driver.py
+1 -1
View File
@@ -230,7 +230,7 @@ def _geom_groups(
key = FunctionKey(types, data_ids, condim)
if types[0] == mujoco.mjtGeom.mjGEOM_HFIELD:
if int(types[0]) == int(mujoco.mjtGeom.mjGEOM_HFIELD):
# add static grid bounds to the grouping key for hfield collisions
geom_rbound_hfield = (
m._impl.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound # pytype: disable=attribute-error
mjx/mujoco/mjx/_src/io.py
+3 -3
View File
@@ -358,8 +358,8 @@ def _put_model_jax(
t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.')
# margin/gap not supported for meshes and height fields
no_margin = {mujoco.mjtGeom.mjGEOM_MESH, mujoco.mjtGeom.mjGEOM_HFIELD}
if no_margin.intersection({t1, t2}):
no_margin = {int(mujoco.mjtGeom.mjGEOM_MESH), int(mujoco.mjtGeom.mjGEOM_HFIELD)}
if no_margin.intersection({int(t1), int(t2)}):
if ip != -1:
margin = m.pair_margin[ip]
else:
@@ -368,7 +368,7 @@ def _put_model_jax(
t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.')
for t, g in [(t1, g1), (t2, g2)]:
if t == mujoco.mjtGeom.mjGEOM_MESH:
if int(t) == int(mujoco.mjtGeom.mjGEOM_MESH):
mesh_geomid.add(g)
for enum_field, enum_type, mj_type in (
mjx/mujoco/mjx/_src/sensor.py
+29 -28
View File
@@ -61,22 +61,22 @@ def sensor_pos(m: Model, d: Data) -> Data:
# position and orientation by object type
objtype_data = {
ObjType.UNKNOWN: (
int(ObjType.UNKNOWN): (
np.zeros((1, 3)),
np.expand_dims(np.eye(3), axis=0),
), # world
ObjType.BODY: (d.xipos, d.ximat),
ObjType.XBODY: (d.xpos, d.xmat),
ObjType.GEOM: (d.geom_xpos, d.geom_xmat),
ObjType.SITE: (d.site_xpos, d.site_xmat),
ObjType.CAMERA: (d.cam_xpos, d.cam_xmat),
int(ObjType.BODY): (d.xipos, d.ximat),
int(ObjType.XBODY): (d.xpos, d.xmat),
int(ObjType.GEOM): (d.geom_xpos, d.geom_xmat),
int(ObjType.SITE): (d.site_xpos, d.site_xmat),
int(ObjType.CAMERA): (d.cam_xpos, d.cam_xmat),
}
# frame axis indexing
frame_axis = {
SensorType.FRAMEXAXIS: 0,
SensorType.FRAMEYAXIS: 1,
SensorType.FRAMEZAXIS: 2,
int(SensorType.FRAMEXAXIS): 0,
int(SensorType.FRAMEYAXIS): 1,
int(SensorType.FRAMEZAXIS): 2,
}
stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS
@@ -190,8 +190,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
for ot, rt in set(zip(objtype, reftype)):
idxt = (objtype == ot) & (reftype == rt)
refidt = refid[idxt]
xpos, _ = objtype_data[ot]
xpos_ref, xmat_ref = objtype_data[rt]
xpos, _ = objtype_data[int(ot)]
xpos_ref, xmat_ref = objtype_data[int(rt)]
xpos = xpos[objid[idxt]]
xpos_ref = xpos_ref[refidt]
xmat_ref = xmat_ref[refidt]
@@ -211,8 +211,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
for ot, rt in set(zip(objtype, reftype)):
idxt = (objtype == ot) & (reftype == rt)
refidt = refid[idxt]
_, xmat = objtype_data[ot]
_, xmat_ref = objtype_data[rt]
_, xmat = objtype_data[int(ot)]
_, xmat_ref = objtype_data[int(rt)]
xmat = xmat[objid[idxt]]
xmat_ref = xmat_ref[refidt]
cutofft = cutoff[idxt]
@@ -224,6 +224,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.FRAMEQUAT:
def _quat(otype, oid):
otype = int(otype)
if otype == ObjType.XBODY:
return d.xquat[oid]
elif otype == ObjType.BODY:
@@ -435,23 +436,23 @@ def sensor_acc(m: Model, d: Data) -> Data:
# position and bodyid by object type
objtype_data = {
ObjType.UNKNOWN: (np.zeros((1, 3)), np.arange(1)),
ObjType.BODY: (d.xipos, np.arange(m.nbody)),
ObjType.XBODY: (d.xpos, np.arange(m.nbody)),
ObjType.GEOM: (d.geom_xpos, m.geom_bodyid),
ObjType.SITE: (d.site_xpos, m.site_bodyid),
ObjType.CAMERA: (d.cam_xpos, m.cam_bodyid),
int(ObjType.UNKNOWN): (np.zeros((1, 3)), np.arange(1)),
int(ObjType.BODY): (d.xipos, np.arange(m.nbody)),
int(ObjType.XBODY): (d.xpos, np.arange(m.nbody)),
int(ObjType.GEOM): (d.geom_xpos, m.geom_bodyid),
int(ObjType.SITE): (d.site_xpos, m.site_bodyid),
int(ObjType.CAMERA): (d.cam_xpos, m.cam_bodyid),
}
stage_acc = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC
sensor_types = set(m.sensor_type[stage_acc])
sensor_types = {int(st) for st in set(m.sensor_type[stage_acc])}
if sensor_types & {
SensorType.ACCELEROMETER,
SensorType.FORCE,
SensorType.TORQUE,
SensorType.FRAMELINACC,
SensorType.FRAMEANGACC,
int(SensorType.ACCELEROMETER),
int(SensorType.FORCE),
int(SensorType.TORQUE),
int(SensorType.FRAMELINACC),
int(SensorType.FRAMEANGACC),
}:
d = smooth.rne_postconstraint(m, d)
@@ -459,8 +460,8 @@ def sensor_acc(m: Model, d: Data) -> Data:
contact_maxforce = (contact_intprm[:, 1] == 2).any()
contact_dataforce = (contact_intprm[:, 0] & (1 << 1)).any()
contact_datatorque = (contact_intprm[:, 0] & (1 << 2)).any()
if (m.sensor_type[stage_acc] == SensorType.TOUCH).any() | (
(m.sensor_type[stage_acc] == SensorType.CONTACT).any()
if (m.sensor_type[stage_acc] == int(SensorType.TOUCH)).any() | (
(m.sensor_type[stage_acc] == int(SensorType.CONTACT)).any()
and (contact_maxforce | contact_dataforce | contact_datatorque)
):
# compute contact forces
@@ -741,7 +742,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
for ot in set(objtype):
idxt = objtype == ot
objidt = objid[idxt]
pos, bodyid = objtype_data[ot]
pos, bodyid = objtype_data[int(ot)]
pos = pos[objidt]
bodyid = bodyid[objidt]
cacc = d._impl.cacc[bodyid] # pyrefly: ignore[missing-attribute]
mjx/mujoco/mjx/_src/support.py
+6 -4
View File
@@ -219,7 +219,7 @@ def local_to_global(
def _getnum(m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj) -> int:
"""Gets the number of objects for the given object type."""
return {
counts = {
mujoco.mjtObj.mjOBJ_BODY: m.nbody,
mujoco.mjtObj.mjOBJ_JOINT: m.njnt,
mujoco.mjtObj.mjOBJ_GEOM: m.ngeom,
@@ -236,14 +236,15 @@ def _getnum(m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj) -> int:
mujoco.mjtObj.mjOBJ_TUPLE: m.ntuple,
mujoco.mjtObj.mjOBJ_KEY: m.nkey,
mujoco.mjtObj.mjOBJ_FLEX: m.nflex,
}.get(obj, 0)
}
return {int(k): v for k, v in counts.items()}.get(int(obj), 0)
def _getadr(
m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj
) -> np.ndarray:
"""Gets the name addresses for the given object type."""
return {
adrs = {
mujoco.mjtObj.mjOBJ_BODY: m.name_bodyadr,
mujoco.mjtObj.mjOBJ_JOINT: m.name_jntadr,
mujoco.mjtObj.mjOBJ_GEOM: m.name_geomadr,
@@ -260,7 +261,8 @@ def _getadr(
mujoco.mjtObj.mjOBJ_TUPLE: m.name_tupleadr,
mujoco.mjtObj.mjOBJ_KEY: m.name_keyadr,
mujoco.mjtObj.mjOBJ_FLEX: m.name_flexadr,
}[obj]
}
return {int(k): v for k, v in adrs.items()}[int(obj)]
def id2name(