From 0e3306b62dbd9aa8f287e70422ea1634e1121c1c Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Fri, 3 Apr 2026 00:12:46 -0700
Subject: [PATCH] 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 | 2 ++
 mjx/mujoco/mjx/_src/types.py   | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py
index 24031821..ee2e0e40 100644
--- a/mjx/mujoco/mjx/_src/support.py
+++ b/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]
 
 
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index f4cda34c..8bc60bd2 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/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