Refactor: Expose flex-related fields directly on mjx.Model.

PiperOrigin-RevId: 896958292
Change-Id: I94b32048e69d2b6844db0911f037318d5fbf4a6a

mjx/mujoco/mjx/warp/bvh.py

@@ -48,7 +48,6 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model
@@ -134,12 +133,12 @@ def _refit_bvh_jax_impl(
       m._impl.flex_radius,
       m._impl.flex_shell,
       m._impl.flex_shelldataadr,
-      m._impl.flex_vertadr,
-      m._impl.flex_vertnum,
+      m.flex_vertadr,
+      m.flex_vertnum,
       m.geom_dataid,
       m.geom_size,
       m.geom_type,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexelem,
       d._impl.flexvert_xpos,
       d.geom_xmat,

mjx/mujoco/mjx/warp/collision_driver.py

@@ -46,7 +46,6 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
@@ -371,7 +370,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_shell,
       m._impl.flex_shelldataadr,
       m._impl.flex_shellnum,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m._impl.flex_vertflexid,
       m.geom_aabb,
       m.geom_conaffinity,
@@ -413,7 +412,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vert,
       m.mesh_vertadr,
       m.mesh_vertnum,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexelem,
       m._impl.nflexshelldata,
       m._impl.nflexvert,

mjx/mujoco/mjx/warp/forward.py

@@ -46,7 +46,6 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _forward_shim(
     # Model
@@ -1499,10 +1498,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_shellnum,
       m._impl.flex_stiffness,
       m._impl.flex_vert,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m._impl.flex_vertbodyid,
       m._impl.flex_vertflexid,
-      m._impl.flex_vertnum,
+      m.flex_vertnum,
       m._impl.flexedge_J_colind,
       m._impl.flexedge_J_rowadr,
       m._impl.flexedge_J_rownnz,
@@ -1596,7 +1595,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.nbranch,
       m.ncam,
       m.neq,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexedge,
       m._impl.nflexelem,
       m._impl.nflexshelldata,
@@ -3444,10 +3443,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_shellnum,
       m._impl.flex_stiffness,
       m._impl.flex_vert,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m._impl.flex_vertbodyid,
       m._impl.flex_vertflexid,
-      m._impl.flex_vertnum,
+      m.flex_vertnum,
       m._impl.flexedge_J_colind,
       m._impl.flexedge_J_rowadr,
       m._impl.flexedge_J_rownnz,
@@ -3542,7 +3541,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nbranch,
       m.ncam,
       m.neq,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexedge,
       m._impl.nflexelem,
       m._impl.nflexshelldata,

mjx/mujoco/mjx/warp/render.py

@@ -48,7 +48,6 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
@@ -164,7 +163,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m.cam_sensorsize,
       m._impl.flex_edge,
       m._impl.flex_radius,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m.geom_dataid,
       m.geom_matid,
       m.geom_rgba,

mjx/mujoco/mjx/warp/types.py

@@ -165,10 +165,8 @@ class ModelWarp(PyTreeNode):
   flex_shellnum: np.ndarray
   flex_stiffness: np.ndarray
   flex_vert: np.ndarray
-  flex_vertadr: np.ndarray
   flex_vertbodyid: np.ndarray
   flex_vertflexid: np.ndarray
-  flex_vertnum: np.ndarray
   flexedge_J_colind: np.ndarray
   flexedge_J_rowadr: np.ndarray
   flexedge_J_rownnz: np.ndarray
@@ -200,7 +198,6 @@ class ModelWarp(PyTreeNode):
   nJfe: int
   nacttrnbody: int
   nbranch: int
-  nflex: int
   nflexedge: int
   nflexelem: int
   nflexelemdata: int