Add useful flex arrays to MJX API.

This enables using name2id for flexes.

PiperOrigin-RevId: 893923892
Change-Id: I7c2a71791e26fa2280cfa75fd66db328f4599b0b

mjx/mujoco/mjx/_src/support.py

@@ -248,6 +248,7 @@ def _getnum(m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj) -> int:
       mujoco.mjtObj.mjOBJ_NUMERIC: m.nnumeric,
       mujoco.mjtObj.mjOBJ_TUPLE: m.ntuple,
       mujoco.mjtObj.mjOBJ_KEY: m.nkey,
+      mujoco.mjtObj.mjOBJ_FLEX: m.nflex,
   }.get(obj, 0)
 
 
@@ -271,6 +272,7 @@ def _getadr(
       mujoco.mjtObj.mjOBJ_NUMERIC: m.name_numericadr,
       mujoco.mjtObj.mjOBJ_TUPLE: m.name_tupleadr,
       mujoco.mjtObj.mjOBJ_KEY: m.name_keyadr,
+      mujoco.mjtObj.mjOBJ_FLEX: m.name_flexadr,
   }[obj]
 
 

mjx/mujoco/mjx/_src/types.py

@@ -611,6 +611,7 @@ class Model(PyTreeNode):
   nsite: int
   ncam: int
   nlight: int
+  nflex: int
   nmesh: int
   nmeshvert: int
   nmeshnormal: int
@@ -774,6 +775,9 @@ class Model(PyTreeNode):
   mesh_texcoordadr: np.ndarray
   mesh_texcoordnum: np.ndarray
   mesh_texcoord: np.ndarray
+  flex_vertadr: np.ndarray
+  flex_vertnum: np.ndarray
+  flex_vert0: np.ndarray
   hfield_size: np.ndarray
   hfield_nrow: np.ndarray
   hfield_ncol: np.ndarray
@@ -881,6 +885,7 @@ class Model(PyTreeNode):
   name_geomadr: np.ndarray
   name_siteadr: np.ndarray
   name_camadr: np.ndarray
+  name_flexadr: np.ndarray
   name_meshadr: np.ndarray
   name_hfieldadr: np.ndarray
   name_pairadr: np.ndarray