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

PiperOrigin-RevId: 966094569
Change-Id: I014226016df9d519f250d0fc25045064cd94cb1e

mjx/mujoco/mjx/_src/collision_driver.py

@@ -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

@@ -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}
+    no_margin = {int(mujoco.mjtGeom.mjGEOM_MESH), int(mujoco.mjtGeom.mjGEOM_HFIELD)}
-    if no_margin.intersection({t1, t2}):
+    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

@@ -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),
+      int(ObjType.BODY): (d.xipos, d.ximat),
-      ObjType.XBODY: (d.xpos, d.xmat),
+      int(ObjType.XBODY): (d.xpos, d.xmat),
-      ObjType.GEOM: (d.geom_xpos, d.geom_xmat),
+      int(ObjType.GEOM): (d.geom_xpos, d.geom_xmat),
-      ObjType.SITE: (d.site_xpos, d.site_xmat),
+      int(ObjType.SITE): (d.site_xpos, d.site_xmat),
-      ObjType.CAMERA: (d.cam_xpos, d.cam_xmat),
+      int(ObjType.CAMERA): (d.cam_xpos, d.cam_xmat),
   }
 
   # frame axis indexing
   frame_axis = {
-      SensorType.FRAMEXAXIS: 0,
+      int(SensorType.FRAMEXAXIS): 0,
-      SensorType.FRAMEYAXIS: 1,
+      int(SensorType.FRAMEYAXIS): 1,
-      SensorType.FRAMEZAXIS: 2,
+      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, _ = objtype_data[int(ot)]
-        xpos_ref, xmat_ref = objtype_data[rt]
+        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 = objtype_data[int(ot)]
-        _, xmat_ref = objtype_data[rt]
+        _, 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)),
+      int(ObjType.UNKNOWN): (np.zeros((1, 3)), np.arange(1)),
-      ObjType.BODY: (d.xipos, np.arange(m.nbody)),
+      int(ObjType.BODY): (d.xipos, np.arange(m.nbody)),
-      ObjType.XBODY: (d.xpos, np.arange(m.nbody)),
+      int(ObjType.XBODY): (d.xpos, np.arange(m.nbody)),
-      ObjType.GEOM: (d.geom_xpos, m.geom_bodyid),
+      int(ObjType.GEOM): (d.geom_xpos, m.geom_bodyid),
-      ObjType.SITE: (d.site_xpos, m.site_bodyid),
+      int(ObjType.SITE): (d.site_xpos, m.site_bodyid),
-      ObjType.CAMERA: (d.cam_xpos, m.cam_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,
+      int(SensorType.ACCELEROMETER),
-      SensorType.FORCE,
+      int(SensorType.FORCE),
-      SensorType.TORQUE,
+      int(SensorType.TORQUE),
-      SensorType.FRAMELINACC,
+      int(SensorType.FRAMELINACC),
-      SensorType.FRAMEANGACC,
+      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() | (
+  if (m.sensor_type[stage_acc] == int(SensorType.TOUCH)).any() | (
-      (m.sensor_type[stage_acc] == SensorType.CONTACT).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

@@ -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(